\n",
" | 1 | \n",
" Flash | \n",
" BF16/FP16 | \n",
@@ -208,11 +200,11 @@
"source": [
"## 2. Backend Selection\n",
"\n",
- "Given the various attention backends, Transformer Engine has a selection logic in place to choose the most appropriate backend for a particular set of user inputs and runtime environment. The selection logic is based on both backend availability and backend performance.\n",
+ "Given the various attention backends, Transformer Engine first determines which backends are eligible for the provided inputs and runtime environment, then applies a preference order among the eligible backends. Eligibility is affected by user environment variables, GPU architecture, installed `flash-attn` and cuDNN versions, data type and FP8 recipe, QKV layout, training or inference mode, dropout, and other attention features.\n",
"\n",
- "Backend availability is determined by factors such as model configuration, training hyper-parameters, software versions, and the GPU architecture in question. For example, some considerations are the sequence length, number of attention heads, head size, attention mask type, attention bias type, training or inference mode, self or cross attention, MHA or MQA/GQA, `flash-attn`/cuDNN library versions, and the compute capability of the GPU.\n",
+ "In PyTorch, the candidates are FlashAttention (`flash-attn` v2, v3, or v4), FusedAttention (cuDNN sub-backends), and UnfusedDotProductAttention. Users can disable whole backend families with `NVTE_FLASH_ATTN`, `NVTE_FUSED_ATTN`, or `NVTE_UNFUSED_ATTN`. In JAX, Transformer Engine checks whether a cuDNN fused-attention kernel is available when `NVTE_FUSED_ATTN=1`; otherwise it falls back to the JAX-native implementation.\n",
"\n",
- "When there are multiple backends available, Transformer Engine makes backend selection based on performance. In general, there are a few rules being followed in our selection logic (see table below). As we monitor the performance of different backends, the selection logic may change.\n",
+ "At a high level, the architecture-specific PyTorch selection order is:\n",
"\n",
"\n",
" \n",
@@ -220,22 +212,29 @@
" | Selection Order | \n",
"
\n",
" \n",
- " | PyTorch | \n",
- " sm90: cuDNN attention > flash-attention > PyTorch-native attention | \n",
+ " PyTorch | \n",
+ " sm8x (Ampere/Ada): flash-attention > cuDNN attention > PyTorch-native attention | \n",
"
\n",
" \n",
- " | sm80: flash-attention > cuDNN attention > PyTorch-native attention | \n",
+ " sm90 (Hopper): cuDNN attention > flash-attention > PyTorch-native attention | \n",
"
\n",
" \n",
- " | \n",
- " cuDNN attention: sub-backend 1 > sub-backend 0\n",
- " | \n",
+ " sm100/sm120 (Blackwell): cuDNN attention > flash-attention > PyTorch-native attention | \n",
+ "
\n",
+ " \n",
+ " | cuDNN attention: BF16/FP16 uses sub-backend 1 when eligible; FP8 uses sub-backend 2 when enabled and eligible | \n",
"
\n",
" \n",
" | JAX | \n",
" cuDNN attention > JAX-native attention | \n",
"
\n",
- "
"
+ "\n",
+ "\n",
+ "Within FlashAttention, TE uses the installed implementation that is supported for the architecture and input. FlashAttention 3 is Hopper-only (`sm90`). FlashAttention 4 supports `sm80`, `sm90`, `sm100`, and `sm120`; on Hopper, TE prefers FlashAttention 3 over FlashAttention 4 when both are installed and eligible. On Blackwell, FlashAttention 4 is the Blackwell-specific flash-attention path when installed and eligible, while FlashAttention 2 can still be eligible depending on the installed version and input configuration.\n",
+ "\n",
+ "Within cuDNN FusedAttention, TE asks the fused-attention helper which sub-backend is eligible. Sub-backend 1 is the BF16/FP16 flash-based path when available; sub-backend 2 is the FP8 path when FP8 DPA is enabled and the architecture, cuDNN version, and input configuration support it. Hopper supports eligible FP8 DPA through cuDNN sub-backend 2. In the current PyTorch selector, eligible FP8 DPA on Blackwell is an `sm100` path and is disabled on `sm120`.\n",
+ "\n",
+ "When all optimized backends are disabled or ineligible, TE falls back to UnfusedDotProductAttention if it is enabled. If no backend is eligible, backend selection returns no backend and the caller raises an error. As we monitor the performance of different backends, the selection logic may change."
]
},
{
@@ -350,7 +349,7 @@
"**cuDNN attention sub-backends:**\n",
"This environment variable allows users to express their preference of cuDNN attention sub-backends. However, the elected sub-backend will only be used *if* it is eligible, i.e. if it has support for the provided inputs and runtime environment.\n",
"```\n",
- "NVTE_FUSED_ATTN_BACKEND = 0/1/2 # user preference of cuDNN sub-backend\n",
+ "NVTE_FUSED_ATTN_BACKEND = 1/2 # user preference of cuDNN sub-backend\n",
"```\n",
"\n",
"**Execution paths of cuDNN sub-backend 1:**\n",
@@ -369,7 +368,7 @@
"\n",
"Note\n",
" \n",
- "Environment variables NVTE_FLASH_ATTN, NVTE_FUSED_ATTN, NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT and NVTE_ALLOW_NONDETERMINISTIC_ALGO are only supported in PyTorch, and will be added to JAX in the future.\n",
+ "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, NVTE_FUSED_ATTN_BACKEND, NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n",
"
\n",
"\n",
"### 2.3 Example Tests\n",
diff --git a/examples/jax/ep/bench/run_ep_bench.sh b/examples/jax/ep/bench/run_ep_bench.sh
index 1531dfd5cf..63133156eb 100755
--- a/examples/jax/ep/bench/run_ep_bench.sh
+++ b/examples/jax/ep/bench/run_ep_bench.sh
@@ -47,6 +47,13 @@ NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
if [ "${NUM_GPUS}" -lt 4 ]; then
echo "EP bench requires >=4 GPUs (found ${NUM_GPUS}); SKIPPING."; exit 0
fi
+
+# NCCL EP requires active NVLink P2P among ranks on the node.
+if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then
+ echo "NVLink not detected on this platform — EP bench requires NVLink; SKIPPING."
+ exit 0
+fi
+
NUM=4
COORD="${COORD:-127.0.0.1:23457}"
TIMEOUT_S="${TIMEOUT_S:-1800}"
diff --git a/examples/pytorch/ep/bench/ep_bench.py b/examples/pytorch/ep/bench/ep_bench.py
new file mode 100644
index 0000000000..2b7a2c62e5
--- /dev/null
+++ b/examples/pytorch/ep/bench/ep_bench.py
@@ -0,0 +1,421 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""PyTorch EP perf bench: raw and autograd dispatch/combine on a single EP group.
+
+One process per GPU; launched via run_ep_bench.sh (torchrun).
+
+Stages (each timed in its own loop):
+ - dispatch_raw: _ep_dispatch_raw (no autograd, no prepare)
+ - ep_dispatch_fwd: ep_dispatch forward only
+ - ep_dispatch_fwd_bwd: ep_dispatch + backward on 0.5 * ||recv||^2
+ - combine_raw: _ep_combine_raw (no autograd)
+ - ep_combine_fwd: ep_combine forward only
+ - ep_combine_fwd_bwd: ep_combine + backward
+
+ep_prepare runs once outside the timed loops. --kineto DIR dumps a Chrome
+trace plus a per-kernel summary on rank 0.
+"""
+
+import argparse
+import gc
+import os
+import sys
+import time
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from transformer_engine.pytorch.ep import (
+ EpBuffer,
+ ep_bootstrap,
+ ep_combine,
+ ep_dispatch,
+ ep_finalize,
+ ep_prepare,
+ _ep_combine_raw,
+ _ep_dispatch_raw,
+)
+
+
+def _parse_args():
+ p = argparse.ArgumentParser(description="TE-PyTorch EP perf bench")
+ p.add_argument("--tokens-per-rank", type=int, default=8192)
+ p.add_argument("--hidden", type=int, default=7168)
+ p.add_argument("--top-k", type=int, default=8)
+ p.add_argument("--num-experts", type=int, default=256)
+ p.add_argument("--warmup", type=int, default=2)
+ p.add_argument("--iters", type=int, default=10)
+ p.add_argument(
+ "--max-num-sms",
+ type=int,
+ default=0,
+ help="Max SMs for dispatch/combine/preprocess kernels (0 = auto).",
+ )
+ p.add_argument(
+ "--kineto",
+ default=None,
+ help="If set, dump a Kineto Chrome trace + per-kernel summary into this dir (rank 0).",
+ )
+ p.add_argument(
+ "--cuda-graph",
+ action="store_true",
+ default=False,
+ help=(
+ "Capture each stage into a CUDA graph and time replay() instead of the eager call. "
+ "Raw + fwd-only stages use torch.cuda.graph; fwd+bwd stages use "
+ "torch.cuda.make_graphed_callables to capture forward and backward together."
+ ),
+ )
+ p.add_argument(
+ "--mode-label",
+ default=None,
+ help="Optional suffix for NVTX range names (e.g. 'fused' / 'unfused').",
+ )
+ p.add_argument(
+ "--caller-provides-dispatch-recv-tokens",
+ action="store_true",
+ default=False,
+ help="Supply recv_tokens to ep_dispatch instead of letting EpBuffer own it.",
+ )
+ p.add_argument(
+ "--caller-provides-grad-expert-out",
+ action="store_true",
+ default=False,
+ help="Supply the combine backward grad buffer to ep_combine.",
+ )
+ return p.parse_args()
+
+
+def _nvtx_funcs():
+ """Return push/pop helpers using torch.cuda.nvtx if available."""
+ try:
+ push = torch.cuda.nvtx.range_push
+ pop = torch.cuda.nvtx.range_pop
+ return push, pop
+ except AttributeError:
+ return lambda _name: None, lambda: None
+
+
+def _device_sm() -> int:
+ major, minor = torch.cuda.get_device_capability()
+ return major * 10 + minor
+
+
+def _make_inputs(rank, world_size, T, H, K, E, device):
+ """Round-robin identity routing + uniform top-k weights."""
+ topk_idx = np.empty((T, K), dtype=np.int64)
+ for t in range(T):
+ for k in range(K):
+ topk_idx[t, k] = ((rank * T + t) * K + k) % E
+ rng = np.random.default_rng(seed=42 + rank)
+ tokens_np = (rng.standard_normal((T, H), dtype=np.float32) * 0.5).astype(np.float32)
+ return (
+ torch.from_numpy(topk_idx).to(device),
+ torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16),
+ torch.full((T, K), 1.0 / K, dtype=torch.float32, device=device),
+ )
+
+
+def _time_stage_us(name, fn, iters, nvtx_suffix, push, pop):
+ """Time fn for iters iterations after one untimed warmup; returns mean us."""
+ # Run iters+1 times; drop the first (autotune outlier) and frame NVTX from iter 1.
+ total_ns = 0
+ counted = 0
+ for i in range(iters + 1):
+ if i == 1:
+ push(f"{name}{nvtx_suffix}")
+ torch.cuda.synchronize()
+ t0 = time.perf_counter_ns()
+ fn()
+ torch.cuda.synchronize()
+ dt = time.perf_counter_ns() - t0
+ if i == 0:
+ continue
+ total_ns += dt
+ counted += 1
+ pop()
+ return total_ns / 1e3 / counted
+
+
+def main():
+ args = _parse_args()
+ dist.init_process_group(backend="nccl")
+ rank = dist.get_rank()
+ world_size = dist.get_world_size()
+ torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank)))
+ device = torch.device("cuda", torch.cuda.current_device())
+
+ if _device_sm() < 90:
+ if rank == 0:
+ print(f"[ep_bench] SKIPPED: EP requires SM>=90 (got SM{_device_sm()})")
+ dist.destroy_process_group()
+ return
+ if world_size < 4:
+ if rank == 0:
+ print(f"[ep_bench] SKIPPED: EP requires >=4 ranks (got {world_size})")
+ dist.destroy_process_group()
+ return
+
+ ep_size = world_size
+ E = args.num_experts
+ assert E % ep_size == 0, f"num_experts ({E}) must be divisible by ep_size ({ep_size})"
+ num_local_experts = E // ep_size
+ T = args.tokens_per_rank
+ H = args.hidden
+ K = args.top_k
+ # Conservative cap: every token could land on every local expert.
+ recv_pr = world_size * T * K // 2
+ if rank == 0:
+ print(
+ f"[ep_bench] world={world_size} ep={ep_size} T={T} H={H} K={K} "
+ f"E={E} (local={num_local_experts}) recv_pr={recv_pr}"
+ + (f" mode={args.mode_label}" if args.mode_label else ""),
+ flush=True,
+ )
+
+ ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl")
+ ep_bootstrap(
+ ep_group,
+ num_experts=E,
+ max_tokens_per_rank=T,
+ recv_capacity_per_rank=recv_pr,
+ hidden_dim=H,
+ max_num_sms=args.max_num_sms,
+ )
+
+ topk_idx, tokens_hbm, topk_w_hbm = _make_inputs(rank, world_size, T, H, K, E, device)
+
+ # Caller-supplied buffers for the autograd ep_dispatch/ep_combine stages
+ # (normal mode -> plain tensors), reused across iters. None when not opted in.
+ caller_recv_tokens = (
+ torch.empty(recv_pr, H, dtype=torch.bfloat16, device=device)
+ if args.caller_provides_dispatch_recv_tokens
+ else None
+ )
+ caller_grad_expert_out = (
+ torch.empty(recv_pr, H, dtype=torch.bfloat16, device=device)
+ if args.caller_provides_grad_expert_out
+ else None
+ )
+
+ buffer = EpBuffer(
+ top_k=K,
+ max_tokens_per_rank=T,
+ recv_capacity_per_rank=recv_pr,
+ hidden_dim=H,
+ num_local_experts=num_local_experts,
+ dispatch_recv_tokens=caller_recv_tokens,
+ combine_grad_expert_out=caller_grad_expert_out,
+ )
+
+ tokens = tokens_hbm
+ topk_w = topk_w_hbm
+ recv_tokens = torch.empty(recv_pr, H, dtype=torch.bfloat16, device=device)
+ recv_w = torch.empty(recv_pr, dtype=torch.float32, device=device)
+
+ # -- Prepare once outside the timed loops ------------------------------
+ ep_prepare(buffer, topk_idx)
+ torch.cuda.synchronize()
+
+ # Pre-dispatch a steady recv_tokens / recv_w so combine stages have valid input.
+ _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w)
+ torch.cuda.synchronize()
+ # fp-equivalent stand-in for an MLP output.
+ expert_out = recv_tokens.clone()
+
+ nvtx_suffix = f"[{args.mode_label}]" if args.mode_label else ""
+ push, pop = _nvtx_funcs()
+
+ # -- Stage closures ----------------------------------------------------
+ # Persistent fwd+bwd inputs (make_graphed_callables needs stable storage).
+ tokens_p = tokens.detach().clone().requires_grad_(True)
+ eo_p = recv_tokens.detach().clone().requires_grad_(True)
+
+ # Stand-in callables; the cuda-graph branch below swaps in graphed versions.
+ fwd_bwd_dispatch_fn = lambda x: ep_dispatch(buffer, x, topk_idx, topk_w)[0] # noqa: E731
+ fwd_bwd_combine_fn = lambda expert_out: ep_combine(buffer, expert_out) # noqa: E731
+
+ def _dispatch_raw():
+ _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w)
+
+ def _combine_raw():
+ out_buf = torch.empty(T, H, dtype=torch.bfloat16, device=device)
+ _ep_combine_raw(buffer, expert_out, out_buf)
+
+ def _ep_dispatch_fwd():
+ ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w)
+
+ def _ep_dispatch_fwd_bwd():
+ tokens_p.grad = None
+ r = fwd_bwd_dispatch_fn(tokens_p)
+ (0.5 * (r * r).sum(dtype=torch.float32)).backward()
+
+ def _ep_combine_fwd():
+ ep_combine(buffer, recv_tokens)
+
+ def _ep_combine_fwd_bwd():
+ eo_p.grad = None
+ out = fwd_bwd_combine_fn(eo_p)
+ (0.5 * (out * out).sum(dtype=torch.float32)).backward()
+
+ stages = [
+ ("dispatch_raw", _dispatch_raw, True),
+ ("ep_dispatch_fwd", _ep_dispatch_fwd, True),
+ ("ep_dispatch_fwd_bwd", _ep_dispatch_fwd_bwd, False),
+ ("combine_raw", _combine_raw, True),
+ ("ep_combine_fwd", _ep_combine_fwd, True),
+ ("ep_combine_fwd_bwd", _ep_combine_fwd_bwd, False),
+ ]
+ # Third tuple element: True = direct torch.cuda.graph capture; False = use
+ # make_graphed_callables (autograd-aware) instead.
+
+ # -- Warmup -----------------------------------------------------------
+ for _ in range(args.warmup):
+ for _name, fn, _capt in stages:
+ fn()
+ torch.cuda.synchronize()
+
+ # -- Optional CUDA-graph capture --------------------------------------
+ # Capture each capturable stage on a side stream and time .replay()
+ # instead of the eager call. Outputs allocated inside the
+ # autograd.Function's forward go through the per-capture private pool
+ # so addresses stay stable across replays.
+ captured_runners = {}
+ if args.cuda_graph:
+ # Graph fwd+bwd of the autograd-wrapped ops via make_graphed_callables.
+ class _DispatchMod(torch.nn.Module):
+ def forward(self, x):
+ return ep_dispatch(buffer, x, topk_idx, topk_w)[0]
+
+ class _CombineMod(torch.nn.Module):
+ def forward(self, expert_out):
+ return ep_combine(buffer, expert_out)
+
+ disp_mod = _DispatchMod().cuda()
+ comb_mod = _CombineMod().cuda()
+ g_disp, g_comb = torch.cuda.make_graphed_callables(
+ (disp_mod, comb_mod),
+ ((tokens_p,), (eo_p,)),
+ )
+ fwd_bwd_dispatch_fn = g_disp
+ fwd_bwd_combine_fn = g_comb
+
+ # Direct torch.cuda.graph capture for raw + fwd-only stages.
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for name, fn, direct_capturable in stages:
+ if not direct_capturable:
+ continue
+ fn() # prime the allocator for stable replay addresses
+ torch.cuda.synchronize()
+ g = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(g):
+ fn()
+ captured_runners[name] = g
+ torch.cuda.current_stream().wait_stream(side)
+ torch.cuda.synchronize()
+
+ # -- Optional Kineto profiling ----------------------------------------
+ kineto_ctx = nullcontext()
+ if args.kineto and rank == 0:
+ os.makedirs(args.kineto, exist_ok=True)
+ kineto_ctx = torch.profiler.profile(
+ activities=[
+ torch.profiler.ProfilerActivity.CPU,
+ torch.profiler.ProfilerActivity.CUDA,
+ ],
+ record_shapes=False,
+ with_stack=False,
+ )
+
+ # -- Timed loops ------------------------------------------------------
+ results = {}
+ with kineto_ctx as prof:
+ for name, fn, _ in stages:
+ runner = fn
+ if name in captured_runners:
+ # Time replay() instead of the eager call.
+ graph = captured_runners[name]
+ runner = graph.replay
+ results[name] = _time_stage_us(name, runner, args.iters, nvtx_suffix, push, pop)
+
+ if rank == 0:
+ label = f" [{args.mode_label}]" if args.mode_label else ""
+ print("", flush=True)
+ print(f"| stage | mean wall (us){label} |", flush=True)
+ print("|----------------------|---------------:|", flush=True)
+ for name in (
+ "dispatch_raw",
+ "ep_dispatch_fwd",
+ "ep_dispatch_fwd_bwd",
+ "combine_raw",
+ "ep_combine_fwd",
+ "ep_combine_fwd_bwd",
+ ):
+ print(f"| {name:20s} | {results[name]:14.1f} |", flush=True)
+ print(
+ "| (dispatch fwd-raw) |"
+ f" {results['ep_dispatch_fwd'] - results['dispatch_raw']:14.1f} |",
+ flush=True,
+ )
+ print(
+ "| (dispatch bwd-fwd) |"
+ f" {results['ep_dispatch_fwd_bwd'] - results['ep_dispatch_fwd']:14.1f} |",
+ flush=True,
+ )
+ print(
+ "| (combine fwd-raw) |"
+ f" {results['ep_combine_fwd'] - results['combine_raw']:14.1f} |",
+ flush=True,
+ )
+ print(
+ "| (combine bwd-fwd) |"
+ f" {results['ep_combine_fwd_bwd'] - results['ep_combine_fwd']:14.1f} |",
+ flush=True,
+ )
+ print("", flush=True)
+
+ if args.kineto and rank == 0 and prof is not None:
+ trace_path = os.path.join(args.kineto, "ep_bench_trace.json")
+ prof.export_chrome_trace(trace_path)
+ print(f"[ep_bench] kineto trace: {trace_path}", flush=True)
+ print(
+ prof.key_averages().table(sort_by="cuda_time_total", row_limit=30),
+ flush=True,
+ )
+ kern_csv = os.path.join(args.kineto, "ep_bench_kernels.csv")
+ with open(kern_csv, "w") as f:
+ f.write("name,cuda_time_us,cpu_time_us,count\n")
+ for evt in prof.key_averages():
+ if evt.device_time_total == 0 and evt.cpu_time_total == 0:
+ continue
+ f.write(f"{evt.key},{evt.device_time_total},{evt.cpu_time_total},{evt.count}\n")
+ print(f"[ep_bench] per-kernel CSV: {kern_csv}", flush=True)
+
+ # Captured CUDA graphs (when --cuda-graph) hold references to NCCL EP
+ # handles and per-pool streams; drop them and sync before ep_finalize,
+ # otherwise the post-finalize dist.barrier can deadlock against pending
+ # graph state.
+ torch.cuda.synchronize()
+ if args.cuda_graph:
+ fwd_bwd_dispatch_fn = None
+ fwd_bwd_combine_fn = None
+ captured_runners.clear()
+ del g_disp, g_comb, disp_mod, comb_mod
+ del tokens_p, eo_p, buffer, recv_tokens, recv_w, tokens, topk_w, expert_out
+ gc.collect()
+ torch.cuda.synchronize()
+ # Release NCCL EP's borrowed comm before torch destroys it.
+ ep_finalize()
+ dist.barrier()
+ dist.destroy_process_group()
+ sys.stdout.flush()
+ sys.stderr.flush()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh
new file mode 100755
index 0000000000..fefecd7fa9
--- /dev/null
+++ b/examples/pytorch/ep/bench/run_ep_bench.sh
@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+#
+# Launcher for examples/pytorch/ep/bench/ep_bench.py.
+# Examples:
+# bash run_ep_bench.sh # plain run, stdout only
+# bash run_ep_bench.sh --cuda-graph # capture + replay each stage as a CUDA graph
+# bash run_ep_bench.sh --kineto # Chrome trace + per-kernel CSV (rank 0)
+# bash run_ep_bench.sh --nsys # nsys profile on rank 0 -> results/pyt_nsys.nsys-rep
+
+set -uo pipefail
+
+NSYS=0; KINETO=0; CGRAPH=0
+for a in "$@"; do
+ case "$a" in
+ --nsys) NSYS=1 ;;
+ --kineto) KINETO=1 ;;
+ --cuda-graph) CGRAPH=1 ;;
+ *) echo "unknown arg: $a" >&2; exit 2 ;;
+ esac
+done
+if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then
+ echo "--nsys and --kineto both attach CUPTI; pick one." >&2; exit 2
+fi
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+RESULTS="${SCRIPT_DIR}/results"
+mkdir -p "${RESULTS}"
+export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
+
+DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
+NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}"
+if [ "${NUM_GPUS}" -lt 4 ]; then
+ echo "EP bench requires >=4 GPUs (found ${NUM_GPUS}); SKIPPING."; exit 0
+fi
+if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi
+
+: "${TIMEOUT_S:=1800}"
+: "${NCCL_EP_JIT_CACHE_DIR:=${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)}"
+export NCCL_EP_JIT_CACHE_DIR
+mkdir -p "${NCCL_EP_JIT_CACHE_DIR}"
+
+EXTRA_ARGS=()
+TAG="pyt"
+[ "${CGRAPH}" -eq 1 ] && EXTRA_ARGS+=(--cuda-graph) && TAG="${TAG}_cg"
+if [ "${KINETO}" -eq 1 ]; then
+ EXTRA_ARGS+=(--kineto "${RESULTS}/kineto_${TAG}")
+fi
+
+EP_BENCH_EXTRA_FLAGS="${EP_BENCH_EXTRA_FLAGS:-}"
+LAUNCH=(torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_GPUS}"
+ "${SCRIPT_DIR}/ep_bench.py" "${EXTRA_ARGS[@]}" ${EP_BENCH_EXTRA_FLAGS})
+
+if [ "${NSYS}" -eq 1 ]; then
+ NSYS_CMD=(nsys profile
+ --output "${RESULTS}/pyt_${TAG}_nsys"
+ --force-overwrite=true
+ --trace=cuda,nvtx
+ --gpu-metrics-devices=none
+ --cuda-um-cpu-page-faults=false
+ --cuda-um-gpu-page-faults=false)
+ echo "[run_ep_bench] launching with nsys (results/${TAG}_nsys.nsys-rep)"
+ timeout --foreground --signal=TERM "${TIMEOUT_S}" "${NSYS_CMD[@]}" "${LAUNCH[@]}"
+ RC=$?
+else
+ timeout --foreground --signal=TERM "${TIMEOUT_S}" "${LAUNCH[@]}"
+ RC=$?
+fi
+exit $RC
diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh
new file mode 100755
index 0000000000..8f6da04a00
--- /dev/null
+++ b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh
@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+#
+# Launcher for the native NCCL EP ``ep_bench`` (baseline for PyTorch comparison).
+# Usage:
+# bash run_nccl_ep_bench.sh # plain run, stdout only
+# bash run_nccl_ep_bench.sh --nsys # nsys → results/nccl_ep_nsys.nsys-rep
+
+set -uo pipefail
+
+NSYS=0
+for a in "$@"; do
+ case "$a" in
+ --nsys) NSYS=1 ;;
+ *) echo "unknown arg: $a" >&2; exit 2 ;;
+ esac
+done
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+RESULTS="${SCRIPT_DIR}/results"
+mkdir -p "${RESULTS}"
+
+BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench"
+LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib"
+[ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; }
+
+NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
+if [ "${NUM_GPUS}" -lt 4 ]; then
+ echo "NCCL EP bench requires >=4 GPUs (found ${NUM_GPUS}); SKIPPING."; exit 0
+fi
+if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi
+
+if [ "${NSYS}" -eq 1 ]; then
+ ITERS=10
+else
+ ITERS=50
+fi
+ARGS=(--algorithm ht --layout em --tokens 2048 --hidden 7168 --top-k 8
+ --experts 256 --warmup 5 --iters "${ITERS}")
+[ "${NSYS}" -eq 1 ] && ARGS+=(--profile) # enables NVTX ranges + cudaProfilerStart/Stop
+
+CMD=(/usr/local/mpi/bin/mpirun --allow-run-as-root --oversubscribe -np "${NUM_GPUS}"
+ -x LD_LIBRARY_PATH="${LIB}:${LD_LIBRARY_PATH:-}"
+ "${BIN}" "${ARGS[@]}")
+
+if [ "${NSYS}" -eq 1 ]; then
+ CMD=(nsys profile
+ --output "${RESULTS}/nccl_ep_nsys"
+ --force-overwrite=true
+ --capture-range=cudaProfilerApi
+ --capture-range-end=stop
+ --trace=cuda,nvtx,osrt
+ "${CMD[@]}")
+fi
+
+[ "${NSYS}" -eq 1 ] && SUFFIX="_nsys" || SUFFIX=""
+LOG="${RESULTS}/stdout_nccl_ep${SUFFIX}.txt"
+"${CMD[@]}" 2>&1 | tee "${LOG}"
+echo "Done. Log: ${LOG}"
diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py
new file mode 100644
index 0000000000..149185f251
--- /dev/null
+++ b/examples/pytorch/ep/ep_moe.py
@@ -0,0 +1,255 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""End-to-end MoE example: dispatch -> batched expert linear -> combine, fwd + bwd.
+
+One process per GPU; launched via run_test_ep.sh (torchrun).
+"""
+
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from transformer_engine.pytorch.ep import (
+ EpBuffer,
+ ep_bootstrap,
+ ep_combine,
+ ep_dispatch,
+ ep_finalize,
+)
+
+
+def _parse_args():
+ p = argparse.ArgumentParser(description="TE-PyTorch EP MoE example (fwd + bwd)")
+ p.add_argument("--num-tokens", type=int, default=8, help="Per-rank token count.")
+ p.add_argument("--top-k", type=int, default=2)
+ p.add_argument("--hidden", type=int, default=32)
+ p.add_argument("--hidden-out", type=int, default=32)
+ p.add_argument("--num-experts", type=int, default=None)
+ p.add_argument("--check", action="store_true", default=True)
+ p.add_argument(
+ "--benchmark",
+ action="store_true",
+ help="Time fwd over HBM buffers.",
+ )
+ p.add_argument("--benchmark-iters", type=int, default=20)
+ p.add_argument("--benchmark-warmup", type=int, default=5)
+ p.add_argument(
+ "--caller-provides-dispatch-recv-tokens",
+ action="store_true",
+ default=False,
+ help="Supply recv_tokens to ep_dispatch instead of letting EpBuffer own it.",
+ )
+ p.add_argument(
+ "--caller-provides-grad-expert-out",
+ action="store_true",
+ default=False,
+ help="Supply the combine backward grad buffer to ep_combine.",
+ )
+ return p.parse_args()
+
+
+def _make_routing(rank, T, K, E, num_local_experts):
+ """Deterministic routing: topk_idx[t, k] = (rank*NLE + t*K + k) % E."""
+ topk_idx = np.empty((T, K), dtype=np.int64)
+ for t in range(T):
+ for k in range(K):
+ topk_idx[t, k] = (rank * num_local_experts + t * K + k) % E
+ return topk_idx
+
+
+def _batched_expert_linear(recv_tokens, kernels, num_local_experts):
+ """Per-expert linear via bmm; ``recv_pr // num_local_experts`` slots per expert."""
+ recv_pr, _H = recv_tokens.shape
+ H_out = kernels.shape[-1]
+ slots_per_expert = recv_pr // num_local_experts
+ grouped = recv_tokens.view(num_local_experts, slots_per_expert, recv_tokens.shape[-1])
+ out = torch.bmm(grouped, kernels.to(grouped.dtype))
+ return out.view(recv_pr, H_out)
+
+
+def _reference_moe(tokens, topk_idx, topk_w, kernels):
+ T, K = topk_idx.shape
+ H_out = kernels.shape[-1]
+ out = np.zeros((T, H_out), dtype=np.float32)
+ for t in range(T):
+ tok = tokens[t].astype(np.float32)
+ for k in range(K):
+ e = int(topk_idx[t, k])
+ out[t] += float(topk_w[t, k]) * (tok @ kernels[e].astype(np.float32))
+ return out
+
+
+def _reference_grad(tokens, topk_idx, topk_w, kernels):
+ T, K = topk_idx.shape
+ H = tokens.shape[-1]
+ ref_out = _reference_moe(tokens, topk_idx, topk_w, kernels)
+ grad = np.zeros((T, H), dtype=np.float32)
+ for t in range(T):
+ mixed = np.zeros_like(kernels[0])
+ for k in range(K):
+ mixed = mixed + float(topk_w[t, k]) * kernels[int(topk_idx[t, k])]
+ grad[t] = ref_out[t] @ mixed.T
+ return ref_out, grad
+
+
+def main():
+ args = _parse_args()
+
+ dist.init_process_group(backend="nccl")
+ rank = dist.get_rank()
+ world_size = dist.get_world_size()
+ torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank)))
+ device = torch.device("cuda", torch.cuda.current_device())
+
+ major, minor = torch.cuda.get_device_capability()
+ if major * 10 + minor < 90:
+ if rank == 0:
+ print(f"[ep_moe] SKIPPED: EP requires SM>=90 (got SM{major}{minor})")
+ dist.destroy_process_group()
+ return
+
+ if world_size < 4:
+ if rank == 0:
+ print(f"[ep_moe] SKIPPED: EP requires >= 4 ranks (got {world_size})")
+ dist.destroy_process_group()
+ return
+
+ ep_size = world_size
+ num_experts = args.num_experts if args.num_experts is not None else world_size
+ assert num_experts % ep_size == 0
+ num_local_experts = num_experts // ep_size
+ T = args.num_tokens
+ recv_pr = ep_size * T * args.top_k
+
+ ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl")
+ ep_bootstrap(
+ ep_group,
+ num_experts=num_experts,
+ max_tokens_per_rank=T,
+ recv_capacity_per_rank=recv_pr,
+ hidden_dim=args.hidden,
+ )
+ try:
+ _run_layer(
+ args, rank, world_size, ep_size, num_experts, num_local_experts, T, recv_pr, device
+ )
+ finally:
+ ep_finalize()
+ dist.destroy_process_group()
+
+
+def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, T, recv_pr, device):
+ rng = np.random.default_rng(seed=42 + rank)
+ tokens_np = (rng.standard_normal((T, args.hidden), dtype=np.float32) * 0.5).astype(np.float32)
+ topk_idx_np = _make_routing(rank, T, args.top_k, num_experts, num_local_experts)
+ w_np = np.full((T, args.top_k), 1.0 / args.top_k, dtype=np.float32)
+ # Same seed across ranks -> identical kernel array everywhere.
+ kr = np.random.default_rng(seed=42)
+ kernels_np = (
+ kr.standard_normal((num_experts, args.hidden, args.hidden_out), dtype=np.float32)
+ * (1.0 / np.sqrt(args.hidden))
+ ).astype(np.float32)
+
+ tokens = (
+ torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16).requires_grad_(True)
+ )
+ topk_idx = torch.from_numpy(topk_idx_np).to(device)
+ topk_w = torch.from_numpy(w_np).to(device)
+ kernels_local = torch.from_numpy(
+ kernels_np[rank * num_local_experts : (rank + 1) * num_local_experts]
+ ).to(device=device, dtype=torch.bfloat16)
+
+ # Caller-supplied buffers (normal mode -> plain tensors), reused across iters.
+ recv_tokens = (
+ torch.empty(recv_pr, args.hidden, dtype=torch.bfloat16, device=device)
+ if args.caller_provides_dispatch_recv_tokens
+ else None
+ )
+ grad_expert_out = (
+ torch.empty(recv_pr, args.hidden, dtype=torch.bfloat16, device=device)
+ if args.caller_provides_grad_expert_out
+ else None
+ )
+
+ buffer = EpBuffer(
+ top_k=args.top_k,
+ max_tokens_per_rank=T,
+ recv_capacity_per_rank=recv_pr,
+ hidden_dim=args.hidden,
+ num_local_experts=num_local_experts,
+ dispatch_recv_tokens=recv_tokens,
+ combine_grad_expert_out=grad_expert_out,
+ )
+
+ recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w)
+ expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts)
+ # Apply per-slot topk weighting before combine.
+ expert_out = expert_out * recv_w_out.unsqueeze(-1).to(expert_out.dtype)
+ out = ep_combine(buffer, expert_out)
+
+ loss = 0.5 * (out.float() ** 2).sum()
+ loss.backward()
+ torch.cuda.synchronize()
+
+ if rank == 0:
+ print(
+ f"[ep_moe] loss={loss.item():.4f} grad_tokens.shape={tuple(tokens.grad.shape)} "
+ f"ep={ep_size} num_experts={num_experts} recv_pr={recv_pr}"
+ )
+
+ if args.benchmark:
+ # Time dispatch + expert + combine over HBM buffers.
+ import time
+
+ torch.cuda.synchronize()
+ dist.barrier()
+ for _ in range(args.benchmark_warmup):
+ rt, rw, _tc = ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w)
+ expert_out = _batched_expert_linear(rt, kernels_local, num_local_experts)
+ expert_out = expert_out * rw.unsqueeze(-1).to(expert_out.dtype)
+ ep_combine(buffer, expert_out)
+ torch.cuda.synchronize()
+ dist.barrier()
+ t0 = time.perf_counter()
+ for _ in range(args.benchmark_iters):
+ rt, rw, _tc = ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w)
+ expert_out = _batched_expert_linear(rt, kernels_local, num_local_experts)
+ expert_out = expert_out * rw.unsqueeze(-1).to(expert_out.dtype)
+ ep_combine(buffer, expert_out)
+ torch.cuda.synchronize()
+ dt_ms = (time.perf_counter() - t0) * 1000.0 / args.benchmark_iters
+ if rank == 0:
+ print(f"[ep_moe --benchmark] HBM: {dt_ms:.3f} ms/iter (iters={args.benchmark_iters})")
+
+ if args.check:
+ # All-gather inputs/outputs/grads for a global reference comparison.
+ global_tokens = [torch.empty_like(tokens) for _ in range(world_size)]
+ global_topk_idx = [torch.empty_like(topk_idx) for _ in range(world_size)]
+ global_topk_w = [torch.empty_like(topk_w) for _ in range(world_size)]
+ global_out = [torch.empty_like(out) for _ in range(world_size)]
+ global_grad = [torch.empty_like(tokens.grad) for _ in range(world_size)]
+ dist.all_gather(global_tokens, tokens.detach())
+ dist.all_gather(global_topk_idx, topk_idx)
+ dist.all_gather(global_topk_w, topk_w)
+ dist.all_gather(global_out, out.detach())
+ dist.all_gather(global_grad, tokens.grad)
+ if rank == 0:
+ all_tokens = torch.cat(global_tokens).float().cpu().numpy()
+ all_idx = torch.cat(global_topk_idx).cpu().numpy()
+ all_w = torch.cat(global_topk_w).cpu().numpy()
+ all_out = torch.cat(global_out).float().cpu().numpy()
+ all_grad = torch.cat(global_grad).float().cpu().numpy()
+ ref_out, ref_grad = _reference_grad(all_tokens, all_idx, all_w, kernels_np)
+ np.testing.assert_allclose(all_out, ref_out, rtol=5e-2, atol=5e-2)
+ np.testing.assert_allclose(all_grad, ref_grad, rtol=5e-2, atol=5e-2)
+ print(f"[ep_moe] --check PASSED (ref_out.sum()={float(ref_out.sum()):.4f})")
+
+
+if __name__ == "__main__":
+ main()
+ sys.exit(0)
diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh
new file mode 100755
index 0000000000..13b41f4cb2
--- /dev/null
+++ b/examples/pytorch/ep/run_test_ep.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+set -uo pipefail
+
+DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
+NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}"
+if [ "${NUM_GPUS}" -lt 4 ]; then
+ echo "EP requires >= 4 GPUs (found ${NUM_GPUS}); SKIPPING."
+ exit 0
+fi
+if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi
+
+: ${TE_PATH:=/opt/transformerengine}
+: ${TEST_TIMEOUT_S:=120}
+
+SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py"
+export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}"
+
+# Stage JIT cubins on tmpfs for fast iteration.
+: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"}
+export NCCL_EP_JIT_CACHE_DIR
+mkdir -p "$NCCL_EP_JIT_CACHE_DIR"
+
+echo "*** Executing ep_moe.py across ${NUM_GPUS} GPUs (timeout=${TEST_TIMEOUT_S}s) ***"
+timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \
+ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_GPUS}" \
+ "${SCRIPT}" --check 2>&1 | tee stdout_ep_moe.txt
+RC=${PIPESTATUS[0]}
+
+RET=0
+if [ "${RC}" -ne 0 ]; then RET=1; fi
+if grep -qE "(^|]:)FAILED|(^|]:)Traceback" stdout_ep_moe.txt; then RET=1; fi
+rm -f stdout_ep_moe.txt
+exit $RET
diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh
index 7eb34a62e4..50a51353d1 100644
--- a/qa/L1_pytorch_distributed_unittest/test.sh
+++ b/qa/L1_pytorch_distributed_unittest/test.sh
@@ -50,6 +50,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_use
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py"
# debug tests
diff --git a/qa/L1_pytorch_mcore_fsdp_integration/test.sh b/qa/L1_pytorch_mcore_fsdp_integration/test.sh
index d63c66f2ea..e08cb8bb98 100644
--- a/qa/L1_pytorch_mcore_fsdp_integration/test.sh
+++ b/qa/L1_pytorch_mcore_fsdp_integration/test.sh
@@ -4,6 +4,15 @@
set -e
+# This test uses the MXFP8 recipe (--fp8-recipe mxfp8), which is only supported
+# on Blackwell (compute capability 10.0) and newer.
+DEVICE_ARCH_RAW=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -n 1)
+DEVICE_ARCH=$(echo "${DEVICE_ARCH_RAW}" | sed 's/[^0-9]//g')
+if [[ -z "${DEVICE_ARCH}" || ${DEVICE_ARCH} -lt 100 ]]; then
+ echo "Skipping L1_pytorch_mcore_fsdp_integration: MXFP8 requires compute capability 10.0+ (Blackwell), detected compute_cap=${DEVICE_ARCH_RAW:-unknown}."
+ exit 0
+fi
+
# Megatron-LM / Megatron-FSDP commit for main branch on Apr. 10, 2026.
# Necessary to support wgrad accumulate fusion and Megatron-FSDP NCCL UBR,
# and fixes decoupled_grad <> DistOpt usage in Megatron-LM.
diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh
index 04fbdf1643..330b254e7d 100644
--- a/qa/L2_jax_distributed_unittest/test.sh
+++ b/qa/L2_jax_distributed_unittest/test.sh
@@ -13,3 +13,6 @@ mkdir -p "$XML_LOG_DIR"
# Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate.
XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
+
+# NCCL EP multi-process suite. The launcher skips when fewer than 4 GPUs or no NVLink is detected.
+TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/multi_process_launch_ep.sh
diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh
index 38cbc8ad3d..f455ec0df3 100644
--- a/qa/L2_jax_unittest/test.sh
+++ b/qa/L2_jax_unittest/test.sh
@@ -28,7 +28,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
: ${XML_LOG_DIR:=/logs}
mkdir -p "$XML_LOG_DIR"
-NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
+NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements"
# Note: mnist intentionally does NOT set --xla_gpu_deterministic_ops because it
diff --git a/setup.py b/setup.py
index b231a5c55d..1eead737d5 100644
--- a/setup.py
+++ b/setup.py
@@ -138,7 +138,12 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
def _discover_nccl_home() -> str:
- """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig."""
+ """Resolve NCCL_HOME, preferring the NCCL the dynamic loader resolves at runtime.
+
+ Probes in order: NCCL_HOME env var, ldconfig cache, well-known prefixes, then a
+ pip-installed nvidia-nccl-cu* wheel. To test a non-default NCCL (e.g. a wheel), set
+ NCCL_HOME and ensure the runtime loader resolves the same lib (e.g. LD_LIBRARY_PATH).
+ """
env_home = os.environ.get("NCCL_HOME")
if env_home:
if (Path(env_home) / "include" / "nccl.h").exists():
@@ -152,28 +157,11 @@ def _discover_nccl_home() -> str:
# Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu).
lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu")
- # pip-installed NCCL (nvidia-nccl-cu* wheel) lives under nvidia/nccl in
- # site-packages and has no top-level include/lib layout.
- try:
- import importlib.util
-
- spec = importlib.util.find_spec("nvidia.nccl")
- if spec is not None and spec.submodule_search_locations:
- pip_root = Path(next(iter(spec.submodule_search_locations)))
- if (pip_root / "include" / "nccl.h").exists() and any(
- (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names
- ):
- return str(pip_root)
- except (ImportError, ValueError):
- pass
-
- for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"):
- p = Path(cand)
- if (p / "include" / "nccl.h").exists() and any(
- (p / sub / name).exists() for sub in lib_subdirs for name in lib_names
- ):
- return str(p)
-
+ # Prefer the NCCL the dynamic loader will actually resolve at runtime so the
+ # EP build links against the same libnccl that gets loaded. libtransformer_engine
+ # carries no NCCL RUNPATH, so the loader uses ldconfig/system paths; building
+ # against a different NCCL (e.g. a pip wheel) causes ABI mismatches. ldconfig is
+ # the ground truth for runtime resolution, so consult it before well-known prefixes.
try:
out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode()
for line in out.splitlines():
@@ -187,6 +175,28 @@ def _discover_nccl_home() -> str:
except (subprocess.CalledProcessError, FileNotFoundError):
pass
+ for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"):
+ p = Path(cand)
+ if (p / "include" / "nccl.h").exists() and any(
+ (p / sub / name).exists() for sub in lib_subdirs for name in lib_names
+ ):
+ return str(p)
+
+ # Fall back to a pip-installed NCCL (nvidia-nccl-cu* wheel) under nvidia/nccl
+ # in site-packages, used only when no system NCCL is present.
+ try:
+ import importlib.util
+
+ spec = importlib.util.find_spec("nvidia.nccl")
+ if spec is not None and spec.submodule_search_locations:
+ pip_root = Path(next(iter(spec.submodule_search_locations)))
+ if (pip_root / "include" / "nccl.h").exists() and any(
+ (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names
+ ):
+ return str(pip_root)
+ except (ImportError, ValueError):
+ pass
+
raise RuntimeError(
"Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix."
)
diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt
index 1c4d86a3a8..832177c637 100644
--- a/tests/cpp/operator/CMakeLists.txt
+++ b/tests/cpp/operator/CMakeLists.txt
@@ -15,8 +15,10 @@ add_executable(test_operator
test_cast_mxfp8_grouped.cu
test_cast_nvfp4_transpose.cu
test_cast_float8blockwise.cu
+ test_cast_float8blockwise_grouped.cu
test_dequantize_mxfp8.cu
test_dequantize_mxfp8_grouped.cu
+ test_dequantize_float8blockwise_grouped.cu
test_dequantize_nvfp4.cu
test_transpose.cu
test_cast_transpose.cu
diff --git a/tests/cpp/operator/test_cast_float8blockwise_grouped.cu b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu
new file mode 100644
index 0000000000..bc9f104e17
--- /dev/null
+++ b/tests/cpp/operator/test_cast_float8blockwise_grouped.cu
@@ -0,0 +1,417 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+#include
+
+#include "../test_common.h"
+
+using namespace transformer_engine;
+using namespace test;
+
+namespace {
+
+enum class ShapeRep { SAME_BOTH_DIMS = 0, VARYING_FIRST_DIM = 1 };
+enum class ScalingDir { ROWWISE = 0, COLWISE = 1, BOTH = 2 };
+enum class BlockDim { ONE_D = 1, TWO_D = 2 };
+
+constexpr size_t kBlock = 128;
+
+inline size_t align4(size_t x) { return ((x + 3) / 4) * 4; }
+
+// Configure split-quantize reference: call non-grouped nvte_quantize_v2 on each tensor slice.
+// Returns flat host buffers for per-tensor outputs and scales (in their per-tensor natural
+// layout) so the test can index them and compare element-wise against the grouped layout.
+struct PerTensorRef {
+ std::vector> output; // per tensor, FP8 raw bytes (R_t * K)
+ std::vector> output_t; // per tensor, FP8 raw bytes (K * R_t)
+ std::vector> scale_inv; // per tensor, layout per non-grouped impl
+ std::vector> scale_inv_t; // per tensor, layout per non-grouped impl
+};
+
+// Per-expert scale layout helpers mirroring the kernel + cuBLAS grouped GEMM
+// expectation. Each expert's scales occupy a contiguous sub-block of the global
+// scale buffer; these compute per-expert padded sizes (in floats) so the test
+// can both size the buffer and compute per-expert base offsets.
+// 1D rowwise : blocks_X * roundup(M_t, 4)
+// 1D colwise : blocks_y_t * roundup(K, 4)
+// 2D rowwise : blocks_y_t * roundup(blocks_X, 4)
+// 2D colwise : blocks_X * roundup(blocks_y_t, 4)
+inline size_t per_expert_scale_floats(BlockDim block_dim, bool columnwise, size_t M_t, size_t K) {
+ constexpr size_t kBlk = 128;
+ const size_t blocks_X = (K + kBlk - 1) / kBlk;
+ const size_t blocks_y = (M_t + kBlk - 1) / kBlk;
+ if (block_dim == BlockDim::ONE_D) {
+ if (!columnwise) return blocks_X * align4(M_t);
+ return blocks_y * align4(K);
+ }
+ // 2D
+ if (!columnwise) return blocks_y * align4(blocks_X);
+ return blocks_X * align4(blocks_y);
+}
+
+// Cumulative per-expert offset (in floats) for tensor `t`.
+inline size_t per_expert_scale_offset(const std::vector& first_dims, size_t t,
+ BlockDim block_dim, bool columnwise, size_t K) {
+ size_t offset = 0;
+ for (size_t i = 0; i < t; ++i) {
+ offset += per_expert_scale_floats(block_dim, columnwise, first_dims[i], K);
+ }
+ return offset;
+}
+
+template
+void perform_test(ShapeRep shape_rep, BlockDim block_dim, ScalingDir dir,
+ const std::vector& first_dims_h, size_t K,
+ bool force_pow_2_scales, float epsilon) {
+ if (getDeviceComputeCapability() < hopperComputeCapability ||
+ getDeviceComputeCapability() >= blackwellComputeCapability) {
+ GTEST_SKIP();
+ }
+
+ DType itype = TypeInfo::dtype;
+ DType otype = TypeInfo::dtype;
+
+ const size_t num_tensors = first_dims_h.size();
+ size_t R_total = 0;
+ for (size_t m : first_dims_h) {
+ ASSERT_EQ(m % kBlock, 0u) << "Per-tensor first dim must be multiple of 128";
+ R_total += m;
+ }
+ ASSERT_EQ(K % 16u, 0u);
+
+ // Host data
+ std::mt19937 gen(0xC0FFEEu);
+ std::uniform_real_distribution dist(-2.0f, 1.0f);
+ std::vector input_h(R_total * K);
+ for (auto& v : input_h) v = static_cast(dist(gen));
+
+ // Tensor offsets (element offsets)
+ std::vector offsets_h(num_tensors + 1, 0);
+ for (size_t t = 0; t < num_tensors; ++t) {
+ offsets_h[t + 1] = offsets_h[t] + static_cast(first_dims_h[t] * K);
+ }
+ std::vector first_dims_i64(num_tensors);
+ for (size_t t = 0; t < num_tensors; ++t) first_dims_i64[t] = static_cast(first_dims_h[t]);
+
+ const bool use_rowwise = (dir == ScalingDir::ROWWISE || dir == ScalingDir::BOTH);
+ const bool use_colwise = (dir == ScalingDir::COLWISE || dir == ScalingDir::BOTH);
+
+ const NVTEScalingMode mode =
+ (block_dim == BlockDim::ONE_D) ? NVTE_BLOCK_SCALING_1D : NVTE_BLOCK_SCALING_2D;
+
+ // Allocate grouped device buffers.
+ InputType* input_d = nullptr;
+ OutputType* output_d = nullptr;
+ OutputType* output_t_d = nullptr;
+ float* scale_inv_d = nullptr;
+ float* scale_inv_t_d = nullptr;
+ int64_t* offsets_d = nullptr;
+ int64_t* first_dims_d = nullptr;
+
+ const size_t blocks_X = (K + kBlock - 1) / kBlock;
+
+ // Grouped scale buffers are sized as the sum of per-expert padded sub-blocks,
+ // matching cuBLAS grouped FP8 block-scaling GEMM's per-expert layout.
+ size_t scale_inv_elems = 0;
+ size_t scale_inv_t_elems = 0;
+ for (size_t t = 0; t < num_tensors; ++t) {
+ scale_inv_elems += per_expert_scale_floats(block_dim, /*columnwise=*/false, first_dims_h[t], K);
+ scale_inv_t_elems +=
+ per_expert_scale_floats(block_dim, /*columnwise=*/true, first_dims_h[t], K);
+ }
+ std::vector scale_inv_shape = {scale_inv_elems};
+ std::vector scale_inv_t_shape = {scale_inv_t_elems};
+
+ const size_t input_bytes = R_total * K * sizeof(InputType);
+ const size_t output_bytes = R_total * K * sizeof(OutputType);
+
+ cudaMalloc(&input_d, input_bytes);
+ cudaMemcpy(input_d, input_h.data(), input_bytes, cudaMemcpyHostToDevice);
+ cudaMalloc(&offsets_d, (num_tensors + 1) * sizeof(int64_t));
+ cudaMemcpy(offsets_d, offsets_h.data(), (num_tensors + 1) * sizeof(int64_t),
+ cudaMemcpyHostToDevice);
+ if (shape_rep == ShapeRep::VARYING_FIRST_DIM) {
+ cudaMalloc(&first_dims_d, num_tensors * sizeof(int64_t));
+ cudaMemcpy(first_dims_d, first_dims_i64.data(), num_tensors * sizeof(int64_t),
+ cudaMemcpyHostToDevice);
+ }
+ if (use_rowwise) {
+ cudaMalloc(&output_d, output_bytes);
+ cudaMemset(output_d, 0, output_bytes);
+ cudaMalloc(&scale_inv_d, scale_inv_elems * sizeof(float));
+ cudaMemset(scale_inv_d, 0, scale_inv_elems * sizeof(float));
+ }
+ if (use_colwise) {
+ cudaMalloc(&output_t_d, output_bytes);
+ cudaMemset(output_t_d, 0, output_bytes);
+ cudaMalloc(&scale_inv_t_d, scale_inv_t_elems * sizeof(float));
+ cudaMemset(scale_inv_t_d, 0, scale_inv_t_elems * sizeof(float));
+ }
+
+ // Build grouped tensors.
+ std::vector logical_shape_vec = {R_total, K};
+ NVTEShape logical_shape = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size());
+
+ NVTEGroupedTensor in_gt = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors,
+ logical_shape);
+ NVTEGroupedTensor out_gt = nvte_create_grouped_tensor(mode, num_tensors, logical_shape);
+
+ NVTEBasicTensor in_data = {input_d, static_cast(itype), logical_shape};
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseData, &in_data, sizeof(in_data));
+
+ NVTEShape offsets_shape;
+ offsets_shape.ndim = 1;
+ offsets_shape.data[0] = num_tensors + 1;
+ NVTEBasicTensor offsets_bt = {offsets_d, kNVTEInt64, offsets_shape};
+ if (shape_rep == ShapeRep::VARYING_FIRST_DIM) {
+ NVTEShape first_dims_shape;
+ first_dims_shape.ndim = 1;
+ first_dims_shape.data[0] = num_tensors;
+ NVTEBasicTensor first_dims_bt = {first_dims_d, kNVTEInt64, first_dims_shape};
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedFirstDims, &first_dims_bt,
+ sizeof(first_dims_bt));
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedFirstDims, &first_dims_bt,
+ sizeof(first_dims_bt));
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedTensorOffsets, &offsets_bt,
+ sizeof(offsets_bt));
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedTensorOffsets, &offsets_bt,
+ sizeof(offsets_bt));
+ }
+
+ if (use_rowwise) {
+ NVTEBasicTensor out_data = {output_d, static_cast(otype), logical_shape};
+ NVTEShape scale_inv_shape_nv = nvte_make_shape(scale_inv_shape.data(), scale_inv_shape.size());
+ NVTEBasicTensor scale_bt = {scale_inv_d, kNVTEFloat32, scale_inv_shape_nv};
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseData, &out_data, sizeof(out_data));
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseScaleInv, &scale_bt, sizeof(scale_bt));
+ }
+ if (use_colwise) {
+ NVTEBasicTensor out_t_data = {output_t_d, static_cast(otype), logical_shape};
+ NVTEShape scale_inv_t_shape_nv = nvte_make_shape(scale_inv_t_shape.data(),
+ scale_inv_t_shape.size());
+ NVTEBasicTensor scale_t_bt = {scale_inv_t_d, kNVTEFloat32, scale_inv_t_shape_nv};
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedColumnwiseData, &out_t_data,
+ sizeof(out_t_data));
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedColumnwiseScaleInv, &scale_t_bt,
+ sizeof(scale_t_bt));
+ }
+
+ // Run grouped quantize.
+ QuantizationConfigWrapper quant_config;
+ quant_config.set_force_pow_2_scales(force_pow_2_scales);
+ quant_config.set_amax_epsilon(epsilon);
+ nvte_group_quantize(in_gt, out_gt, quant_config, 0);
+ cudaDeviceSynchronize();
+ ASSERT_EQ(cudaGetLastError(), cudaSuccess);
+
+ // Pull grouped outputs back to host.
+ std::vector output_h(use_rowwise ? R_total * K : 0);
+ std::vector output_t_h(use_colwise ? R_total * K : 0);
+ std::vector scale_inv_h(use_rowwise ? scale_inv_elems : 0);
+ std::vector scale_inv_t_h(use_colwise ? scale_inv_t_elems : 0);
+ if (use_rowwise) {
+ cudaMemcpy(output_h.data(), output_d, R_total * K, cudaMemcpyDeviceToHost);
+ cudaMemcpy(scale_inv_h.data(), scale_inv_d, scale_inv_elems * sizeof(float),
+ cudaMemcpyDeviceToHost);
+ }
+ if (use_colwise) {
+ cudaMemcpy(output_t_h.data(), output_t_d, R_total * K, cudaMemcpyDeviceToHost);
+ cudaMemcpy(scale_inv_t_h.data(), scale_inv_t_d, scale_inv_t_elems * sizeof(float),
+ cudaMemcpyDeviceToHost);
+ }
+
+ // Run split-quantize reference per tensor and compare element-wise.
+ for (size_t t = 0; t < num_tensors; ++t) {
+ const size_t M = first_dims_h[t];
+ const size_t row_offset = static_cast(offsets_h[t]) / K;
+
+ std::vector tshape = {M, K};
+ Tensor ref_in("ref_in_" + std::to_string(t), tshape, itype);
+ // The non-grouped 2D kernel requires rowwise output to be allocated even when only colwise
+ // data is consumed. We always allocate both and compare only what the grouped kernel produced.
+ const bool ref_rowwise = (block_dim == BlockDim::TWO_D) ? true : use_rowwise;
+ const bool ref_colwise = use_colwise;
+ Tensor ref_out("ref_out_" + std::to_string(t), tshape, otype, ref_rowwise, ref_colwise, mode);
+
+ // Copy this tensor's input slice into ref_in.
+ {
+ auto* dst = ref_in.rowwise_dptr();
+ const InputType* src = reinterpret_cast(input_d) + row_offset * K;
+ cudaMemcpy(dst, src, M * K * sizeof(InputType), cudaMemcpyDeviceToDevice);
+ }
+
+ QuantizationConfigWrapper qc;
+ qc.set_force_pow_2_scales(force_pow_2_scales);
+ qc.set_amax_epsilon(epsilon);
+ nvte_quantize_v2(ref_in.data(), ref_out.data(), qc, 0);
+ cudaDeviceSynchronize();
+ ASSERT_EQ(cudaGetLastError(), cudaSuccess);
+ ref_out.to_cpu(); // sync output and scale_inv buffers from GPU to CPU
+
+ // Compare data.
+ if (use_rowwise) {
+ const OutputType* ref_data = ref_out.rowwise_cpu_dptr();
+ for (size_t r = 0; r < M; ++r) {
+ for (size_t c = 0; c < K; ++c) {
+ const uint8_t got = output_h[(row_offset + r) * K + c];
+ const uint8_t exp = reinterpret_cast(ref_data)[r * K + c];
+ ASSERT_EQ(got, exp) << "rowwise data mismatch t=" << t << " r=" << r << " c=" << c;
+ }
+ }
+ }
+ if (use_colwise) {
+ const OutputType* ref_data_t = ref_out.columnwise_cpu_dptr();
+ // Per-expert columnwise data: contiguous (K, M_t) block at element offset
+ // K * row_offset, matching cuBLAS grouped GEMM's per-expert data pointer.
+ const size_t expert_data_off = static_cast(row_offset) * K;
+ for (size_t c = 0; c < K; ++c) {
+ for (size_t r = 0; r < M; ++r) {
+ const uint8_t got = output_t_h[expert_data_off + c * M + r];
+ const uint8_t exp = reinterpret_cast(ref_data_t)[c * M + r];
+ ASSERT_EQ(got, exp) << "colwise data mismatch t=" << t << " c=" << c << " r=" << r;
+ }
+ }
+ }
+
+ // Compare scales. Per-expert layout: each expert's scales live in a
+ // contiguous sub-block at per_expert_scale_offset(...).
+ if (block_dim == BlockDim::ONE_D) {
+ const size_t M_pad = align4(M);
+ const size_t K_pad = align4(K);
+ const size_t blocks_y_per_tensor = M / kBlock;
+ if (use_rowwise) {
+ const float* ref_sc = ref_out.rowwise_cpu_scale_inv_ptr();
+ // Per-expert RW: (blocks_X, roundup(M_t, 4)).
+ const size_t expert_off =
+ per_expert_scale_offset(first_dims_h, t, block_dim, false, K);
+ for (size_t bx = 0; bx < blocks_X; ++bx) {
+ for (size_t r = 0; r < M; ++r) {
+ const float got = scale_inv_h[expert_off + bx * M_pad + r];
+ const float exp = ref_sc[bx * M_pad + r];
+ ASSERT_EQ(got, exp) << "1D rowwise scale mismatch t=" << t << " bx=" << bx
+ << " r=" << r;
+ }
+ }
+ }
+ if (use_colwise) {
+ const float* ref_sct = ref_out.columnwise_cpu_scale_inv_ptr();
+ // Per-expert CW: (blocks_y_t, roundup(K, 4)) compact.
+ const size_t expert_off = per_expert_scale_offset(first_dims_h, t, block_dim, true, K);
+ for (size_t by = 0; by < blocks_y_per_tensor; ++by) {
+ for (size_t c = 0; c < K; ++c) {
+ const float got = scale_inv_t_h[expert_off + by * K_pad + c];
+ const float exp = ref_sct[by * K_pad + c];
+ ASSERT_EQ(got, exp) << "1D colwise scale mismatch t=" << t << " by=" << by
+ << " c=" << c;
+ }
+ }
+ }
+ } else {
+ // 2D per-expert: rowwise (blocks_y_t, roundup(blocks_X, 4)); colwise
+ // (blocks_X, roundup(blocks_y_t, 4)) compact.
+ const size_t blocks_y_per_tensor = M / kBlock;
+ const size_t bx_pad = align4(blocks_X);
+ const size_t by_pad_t = align4(blocks_y_per_tensor);
+ if (use_rowwise) {
+ const float* ref_sc = ref_out.rowwise_cpu_scale_inv_ptr();
+ const size_t expert_off =
+ per_expert_scale_offset(first_dims_h, t, block_dim, false, K);
+ for (size_t by = 0; by < blocks_y_per_tensor; ++by) {
+ for (size_t bx = 0; bx < blocks_X; ++bx) {
+ const float got = scale_inv_h[expert_off + by * bx_pad + bx];
+ const float exp = ref_sc[by * bx_pad + bx];
+ ASSERT_EQ(got, exp) << "2D rowwise scale mismatch t=" << t << " by=" << by
+ << " bx=" << bx;
+ }
+ }
+ }
+ if (use_colwise) {
+ const float* ref_sct = ref_out.columnwise_cpu_scale_inv_ptr();
+ const size_t expert_off = per_expert_scale_offset(first_dims_h, t, block_dim, true, K);
+ for (size_t bx = 0; bx < blocks_X; ++bx) {
+ for (size_t by = 0; by < blocks_y_per_tensor; ++by) {
+ const float got = scale_inv_t_h[expert_off + bx * by_pad_t + by];
+ const float exp = ref_sct[bx * by_pad_t + by];
+ ASSERT_EQ(got, exp) << "2D colwise scale mismatch t=" << t << " bx=" << bx
+ << " by=" << by;
+ }
+ }
+ }
+ }
+ }
+
+ nvte_destroy_grouped_tensor(in_gt);
+ nvte_destroy_grouped_tensor(out_gt);
+ cudaFree(input_d);
+ if (output_d) cudaFree(output_d);
+ if (output_t_d) cudaFree(output_t_d);
+ if (scale_inv_d) cudaFree(scale_inv_d);
+ if (scale_inv_t_d) cudaFree(scale_inv_t_d);
+ cudaFree(offsets_d);
+ if (first_dims_d) cudaFree(first_dims_d);
+}
+
+struct TestConfig {
+ ShapeRep shape_rep;
+ BlockDim block_dim;
+ ScalingDir dir;
+ std::vector first_dims;
+ size_t K;
+};
+
+class GroupedFP8BlockwiseTestSuite : public ::testing::TestWithParam {};
+
+TEST_P(GroupedFP8BlockwiseTestSuite, Test) {
+ const TestConfig& cfg = GetParam();
+ perform_test(cfg.shape_rep, cfg.block_dim, cfg.dir, cfg.first_dims, cfg.K,
+ /*force_pow_2_scales=*/false, /*epsilon=*/0.0f);
+}
+
+std::vector make_configs() {
+ std::vector configs;
+ std::vector> uniform = {{128, 128}, {256, 256, 256, 256}};
+ std::vector> jagged = {
+ {128, 256, 384, 512}, {256, 128, 512, 384, 1024}};
+ std::vector Ks = {128, 256, 512};
+ for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) {
+ for (auto dir : {ScalingDir::ROWWISE, ScalingDir::COLWISE, ScalingDir::BOTH}) {
+ for (size_t K : Ks) {
+ for (const auto& v : uniform) {
+ configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, dir, v, K});
+ }
+ for (const auto& v : jagged) {
+ configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, dir, v, K});
+ }
+ }
+ }
+ }
+ return configs;
+}
+
+std::string make_name(const ::testing::TestParamInfo& info) {
+ const auto& c = info.param;
+ std::string s = (c.shape_rep == ShapeRep::SAME_BOTH_DIMS ? "SAME" : "VARYFIRST");
+ s += "_BD" + std::to_string(static_cast(c.block_dim));
+ s += (c.dir == ScalingDir::ROWWISE ? "_RW"
+ : c.dir == ScalingDir::COLWISE ? "_CW" : "_BOTH");
+ s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size());
+ s += "_M";
+ for (size_t m : c.first_dims) s += "_" + std::to_string(m);
+ return s;
+}
+
+INSTANTIATE_TEST_SUITE_P(GroupedFP8Blockwise, GroupedFP8BlockwiseTestSuite,
+ ::testing::ValuesIn(make_configs()), make_name);
+
+} // namespace
diff --git a/tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu b/tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu
new file mode 100644
index 0000000000..48b0a54c9b
--- /dev/null
+++ b/tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu
@@ -0,0 +1,317 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include "../test_common.h"
+
+using namespace transformer_engine;
+using namespace test;
+
+namespace {
+
+enum class ShapeRep { SAME_BOTH_DIMS = 0, VARYING_FIRST_DIM = 1 };
+enum class ScalingDir { ROWWISE = 0, COLWISE = 1 };
+enum class BlockDim { ONE_D = 1, TWO_D = 2 };
+
+constexpr size_t kBlock = 128;
+
+inline size_t align4(size_t x) { return ((x + 3) / 4) * 4; }
+
+// Per-expert padded scale size (in floats), matching the grouped FP8 block-scaling layout that
+// the grouped quantize kernel writes and cuBLAS grouped GEMM consumes:
+// 1D rowwise : blocks_X * roundup(M_t, 4) (scale shape {blocks_X, roundup(M_t, 4)})
+// 1D colwise : blocks_y_t * roundup(K, 4) (scale shape {blocks_y_t, roundup(K, 4)})
+// 2D rowwise : blocks_y_t * roundup(blocks_X,4) (scale shape {blocks_y_t, roundup(blocks_X, 4)})
+// 2D colwise : blocks_X * roundup(blocks_y_t,4)(scale shape {blocks_X, roundup(blocks_y_t,4)})
+inline void per_expert_scale_shape(BlockDim block_dim, bool columnwise, size_t M_t, size_t K,
+ size_t& scale_y, size_t& scale_x) {
+ const size_t blocks_X = (K + kBlock - 1) / kBlock;
+ const size_t blocks_y = (M_t + kBlock - 1) / kBlock;
+ if (block_dim == BlockDim::ONE_D) {
+ if (!columnwise) {
+ scale_y = blocks_X;
+ scale_x = align4(M_t);
+ } else {
+ scale_y = blocks_y;
+ scale_x = align4(K);
+ }
+ } else {
+ if (!columnwise) {
+ scale_y = blocks_y;
+ scale_x = align4(blocks_X);
+ } else {
+ scale_y = blocks_X;
+ scale_x = align4(blocks_y);
+ }
+ }
+}
+
+inline size_t per_expert_scale_floats(BlockDim block_dim, bool columnwise, size_t M_t, size_t K) {
+ size_t y, x;
+ per_expert_scale_shape(block_dim, columnwise, M_t, K, y, x);
+ return y * x;
+}
+
+// Grouped FP8 block-scaling dequantize test.
+//
+// Methodology mirrors test_dequantize_mxfp8_grouped.cu: the grouped dequantize kernel is
+// validated against single-tensor nvte_dequantize called in a loop for each tensor; results must
+// be bitwise identical. We generate random FP8 data and random FP32 scales laid out in the
+// grouped per-expert format, run nvte_group_dequantize, and for each expert slice out its data +
+// scale sub-block, feed it to a per-tensor (direction-only) dequantize, and compare.
+template
+void performTest(ShapeRep shape_rep, BlockDim block_dim, bool rowwise,
+ const std::vector& first_dims_h, size_t K) {
+ // FP8 block-scaling grouped kernels are Hopper-only (SM90-SM99).
+ if (getDeviceComputeCapability() < hopperComputeCapability ||
+ getDeviceComputeCapability() >= blackwellComputeCapability) {
+ GTEST_SKIP();
+ }
+
+ const DType itype = TypeInfo::dtype;
+ const DType otype = TypeInfo::dtype;
+ const bool columnwise = !rowwise;
+
+ const size_t num_tensors = first_dims_h.size();
+ size_t R_total = 0;
+ for (size_t m : first_dims_h) {
+ ASSERT_EQ(m % kBlock, 0u) << "Per-tensor first dim must be a multiple of 128";
+ R_total += m;
+ }
+ ASSERT_EQ(K % 16u, 0u);
+
+ const NVTEScalingMode mode =
+ (block_dim == BlockDim::ONE_D) ? NVTE_BLOCK_SCALING_1D : NVTE_BLOCK_SCALING_2D;
+
+ // Element offsets (both data and, for columnwise, the transposed (K, M_t) block are contiguous
+ // per expert at element offset row_offset * K).
+ std::vector offsets_h(num_tensors + 1, 0);
+ for (size_t t = 0; t < num_tensors; ++t)
+ offsets_h[t + 1] = offsets_h[t] + static_cast(first_dims_h[t] * K);
+ std::vector first_dims_i64(num_tensors);
+ for (size_t t = 0; t < num_tensors; ++t)
+ first_dims_i64[t] = static_cast(first_dims_h[t]);
+
+ // Per-expert scale sub-block offsets (in floats).
+ std::vector scale_off(num_tensors + 1, 0);
+ for (size_t t = 0; t < num_tensors; ++t)
+ scale_off[t + 1] =
+ scale_off[t] + per_expert_scale_floats(block_dim, columnwise, first_dims_h[t], K);
+ const size_t total_scales = scale_off[num_tensors];
+
+ // ---- Random FP8 data (valid normals) + random FP32 scales ----
+ std::mt19937 gen(0xD3C0DEu);
+ const double minAbs = Numeric_Traits::minNorm;
+ const double maxAbs = Numeric_Traits::maxNorm;
+ std::uniform_real_distribution<> dis(minAbs, maxAbs);
+ std::uniform_real_distribution<> dis_sign(-1.0, 1.0);
+ std::uniform_real_distribution scale_dis(0.25f, 4.0f);
+
+ std::vector data_h(R_total * K);
+ for (auto& v : data_h) {
+ double val = dis(gen);
+ if (dis_sign(gen) < 0.0) val = -val;
+ v = static_cast(val);
+ }
+ std::vector scales_h(total_scales);
+ for (auto& s : scales_h) s = scale_dis(gen);
+
+ // ---- Device buffers ----
+ InputType* data_d = nullptr;
+ float* scales_d = nullptr;
+ OutputType* out_grouped_d = nullptr;
+ int64_t* offsets_d = nullptr;
+ int64_t* first_dims_d = nullptr;
+
+ cudaMalloc(&data_d, R_total * K * sizeof(InputType));
+ cudaMemcpy(data_d, data_h.data(), R_total * K * sizeof(InputType), cudaMemcpyHostToDevice);
+ cudaMalloc(&scales_d, total_scales * sizeof(float));
+ cudaMemcpy(scales_d, scales_h.data(), total_scales * sizeof(float), cudaMemcpyHostToDevice);
+ cudaMalloc(&out_grouped_d, R_total * K * sizeof(OutputType));
+ cudaMemset(out_grouped_d, 0, R_total * K * sizeof(OutputType));
+ cudaMalloc(&offsets_d, (num_tensors + 1) * sizeof(int64_t));
+ cudaMemcpy(offsets_d, offsets_h.data(), (num_tensors + 1) * sizeof(int64_t),
+ cudaMemcpyHostToDevice);
+ if (shape_rep == ShapeRep::VARYING_FIRST_DIM) {
+ cudaMalloc(&first_dims_d, num_tensors * sizeof(int64_t));
+ cudaMemcpy(first_dims_d, first_dims_i64.data(), num_tensors * sizeof(int64_t),
+ cudaMemcpyHostToDevice);
+ }
+
+ // ---- Build grouped input (quantized) + output (high precision) tensors ----
+ std::vector logical_shape_vec = {R_total, K};
+ NVTEShape logical_shape = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size());
+ std::vector data_1d = {R_total * K};
+ NVTEShape data_shape = nvte_make_shape(data_1d.data(), data_1d.size());
+ std::vector scale_1d = {total_scales};
+ NVTEShape scale_shape = nvte_make_shape(scale_1d.data(), scale_1d.size());
+
+ NVTEShape offsets_shape;
+ offsets_shape.ndim = 1;
+ offsets_shape.data[0] = num_tensors + 1;
+ NVTEShape first_dims_shape;
+ first_dims_shape.ndim = 1;
+ first_dims_shape.data[0] = num_tensors;
+ NVTEBasicTensor offsets_bt = {offsets_d, kNVTEInt64, offsets_shape};
+ NVTEBasicTensor first_dims_bt = {first_dims_d, kNVTEInt64, first_dims_shape};
+ auto set_shape_meta = [&](NVTEGroupedTensor gt) {
+ if (shape_rep == ShapeRep::VARYING_FIRST_DIM) {
+ nvte_set_grouped_tensor_param(gt, kNVTEGroupedFirstDims, &first_dims_bt,
+ sizeof(first_dims_bt));
+ nvte_set_grouped_tensor_param(gt, kNVTEGroupedTensorOffsets, &offsets_bt, sizeof(offsets_bt));
+ }
+ };
+
+ NVTEGroupedTensor in_gt = nvte_create_grouped_tensor(mode, num_tensors, logical_shape);
+ NVTEBasicTensor in_data_bt = {data_d, static_cast(itype), data_shape};
+ NVTEBasicTensor in_scale_bt = {scales_d, kNVTEFloat32, scale_shape};
+ if (rowwise) {
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseData, &in_data_bt, sizeof(in_data_bt));
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseScaleInv, &in_scale_bt,
+ sizeof(in_scale_bt));
+ } else {
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedColumnwiseData, &in_data_bt,
+ sizeof(in_data_bt));
+ nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedColumnwiseScaleInv, &in_scale_bt,
+ sizeof(in_scale_bt));
+ }
+ set_shape_meta(in_gt);
+
+ NVTEGroupedTensor out_gt =
+ nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape);
+ NVTEBasicTensor out_data_bt = {out_grouped_d, static_cast(otype), data_shape};
+ nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseData, &out_data_bt, sizeof(out_data_bt));
+ set_shape_meta(out_gt);
+
+ // ---- Grouped dequantize ----
+ nvte_group_dequantize(in_gt, out_gt, 0);
+ cudaDeviceSynchronize();
+ {
+ auto err = cudaGetLastError();
+ ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err);
+ }
+ std::vector out_grouped_h(R_total * K);
+ cudaMemcpy(out_grouped_h.data(), out_grouped_d, R_total * K * sizeof(OutputType),
+ cudaMemcpyDeviceToHost);
+
+ // ---- Reference: host dequantize (out = float(fp8) * scale_inv) ----
+ //
+ // There is no non-grouped FP8 block-scaling dequantize to loop over (only the grouped path
+ // implements it), so the reference is computed on the host. The per-expert scale sub-block
+ // indexing mirrors what test_cast_float8blockwise_grouped.cu validates the grouped quantize
+ // kernel writes; this makes the host reference an independent check of the grouped dequant.
+ // 1D rowwise : scale[bx * roundup(M,4) + r] (bx = c/128)
+ // 1D colwise : scale[by * roundup(K,4) + c] (by = r/128)
+ // 2D rowwise : scale[by * roundup(blocks_X,4) + bx]
+ // 2D colwise : scale[bx * roundup(blocks_y,4) + by]
+ // Data is contiguous (M,K) for rowwise and transposed (K,M) for columnwise.
+ auto scale_index = [&](size_t r, size_t c, size_t M) -> size_t {
+ const size_t blocks_X = (K + kBlock - 1) / kBlock;
+ const size_t blocks_y = (M + kBlock - 1) / kBlock;
+ const size_t bx = c / kBlock;
+ const size_t by = r / kBlock;
+ if (block_dim == BlockDim::ONE_D) {
+ return columnwise ? (by * align4(K) + c) : (bx * align4(M) + r);
+ }
+ return columnwise ? (bx * align4(blocks_y) + by) : (by * align4(blocks_X) + bx);
+ };
+
+ for (size_t t = 0; t < num_tensors; ++t) {
+ const size_t M = first_dims_h[t];
+ const size_t row_offset = static_cast(offsets_h[t]) / K;
+ const size_t data_off = row_offset * K; // rowwise (M,K) or colwise transposed (K,M)
+ const size_t s_off = scale_off[t];
+ for (size_t r = 0; r < M; ++r) {
+ for (size_t c = 0; c < K; ++c) {
+ const size_t d_idx = columnwise ? (c * M + r) : (r * K + c);
+ const float fp8v = static_cast(data_h[data_off + d_idx]);
+ const float sc = scales_h[s_off + scale_index(r, c, M)];
+ const float ref = fp8v * sc;
+ const float got = static_cast(out_grouped_h[(row_offset + r) * K + c]);
+ const float rel = std::fabs(got - ref) / std::max(std::fabs(ref), 1e-3f);
+ ASSERT_LT(rel, 1e-2f) << "dequant mismatch t=" << t << " r=" << r << " c=" << c
+ << " got=" << got << " ref=" << ref << " (fp8=" << fp8v
+ << " scale=" << sc << ")";
+ }
+ }
+ }
+
+ nvte_destroy_grouped_tensor(in_gt);
+ nvte_destroy_grouped_tensor(out_gt);
+ cudaFree(data_d);
+ cudaFree(scales_d);
+ cudaFree(out_grouped_d);
+ cudaFree(offsets_d);
+ if (first_dims_d) cudaFree(first_dims_d);
+}
+
+struct TestConfig {
+ ShapeRep shape_rep;
+ BlockDim block_dim;
+ bool rowwise;
+ std::vector first_dims;
+ size_t K;
+};
+
+std::vector make_configs() {
+ std::vector configs;
+ std::vector> uniform = {{128, 128}, {256, 256, 256, 256}};
+ std::vector> jagged = {{128, 256, 384, 512}, {256, 128, 512, 384, 1024}};
+ std::vector Ks = {128, 256, 512};
+ for (auto bd : {BlockDim::ONE_D, BlockDim::TWO_D}) {
+ for (bool rowwise : {true, false}) {
+ for (size_t K : Ks) {
+ for (const auto& v : uniform)
+ configs.push_back({ShapeRep::SAME_BOTH_DIMS, bd, rowwise, v, K});
+ for (const auto& v : jagged)
+ configs.push_back({ShapeRep::VARYING_FIRST_DIM, bd, rowwise, v, K});
+ }
+ }
+ }
+ return configs;
+}
+
+} // namespace
+
+class GroupedDequantizeFP8BlockwiseTestSuite
+ : public ::testing::TestWithParam> {};
+
+TEST_P(GroupedDequantizeFP8BlockwiseTestSuite, Test) {
+ const TestConfig cfg = std::get<0>(GetParam());
+ const DType output_type = std::get<1>(GetParam());
+ // FP8 block scaling is E4M3-centric (matches the grouped quantize test scope).
+ TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(
+ output_type, OutputType,
+ performTest(cfg.shape_rep, cfg.block_dim, cfg.rowwise, cfg.first_dims,
+ cfg.K););
+}
+
+INSTANTIATE_TEST_SUITE_P(
+ GroupedFP8Blockwise, GroupedDequantizeFP8BlockwiseTestSuite,
+ ::testing::Combine(::testing::ValuesIn(make_configs()),
+ ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)),
+ [](const testing::TestParamInfo& info) {
+ const TestConfig& c = std::get<0>(info.param);
+ std::string s = (c.shape_rep == ShapeRep::SAME_BOTH_DIMS ? "SAME" : "VARYFIRST");
+ s += "_BD" + std::to_string(static_cast(c.block_dim));
+ s += (c.rowwise ? "_RW" : "_CW");
+ s += "_K" + std::to_string(c.K) + "_N" + std::to_string(c.first_dims.size());
+ s += "_" + test::typeName(std::get<1>(info.param));
+ return s;
+ });
diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu
index 12b4703469..09fcbad8df 100644
--- a/tests/cpp/operator/test_grouped_gemm.cu
+++ b/tests/cpp/operator/test_grouped_gemm.cu
@@ -385,6 +385,36 @@ inline AlphaBetaTensors make_alpha_beta(size_t num_gemms) {
// Compare each tensor inside a grouped D buffer (with per-tensor offsets) against the
// reference D_multi[i] tensors.
+// Capture `body(stream)` into a CUDA graph and replay it `n_replays` times, calling
+// `verify(iter)` after each launch+sync. Used to assert that the grouped GEMM kernels
+// are graph-safe (capture succeeds, replay is correct, and replays are deterministic).
+template
+inline void capture_and_replay_grouped_gemm(cudaStream_t stream, int n_replays,
+ Body&& body, Verify&& verify) {
+ // Warmup off-graph so cuBLASLt can initialize its internal heuristic / handle state
+ // outside capture (the very first matmul on a handle may do host-side setup that
+ // isn't capturable in some cuBLAS versions).
+ body(stream);
+ NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+
+ NVTE_CHECK_CUDA(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed));
+ body(stream);
+ cudaGraph_t graph;
+ NVTE_CHECK_CUDA(cudaStreamEndCapture(stream, &graph));
+
+ cudaGraphExec_t exec;
+ NVTE_CHECK_CUDA(cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0));
+
+ for (int i = 0; i < n_replays; ++i) {
+ NVTE_CHECK_CUDA(cudaGraphLaunch(exec, stream));
+ NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+ verify(i);
+ }
+
+ NVTE_CHECK_CUDA(cudaGraphExecDestroy(exec));
+ NVTE_CHECK_CUDA(cudaGraphDestroy(graph));
+}
+
inline void compare_grouped_d_to_multi(
const GroupedBuffers& grouped_D,
const std::vector>& shapes,
@@ -649,6 +679,255 @@ void run_grouped_gemm_discrete_in_case(const TestParams& params) {
compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, "grouped_discrete_in_vs_multi");
}
+// Graph-capture variant of run_grouped_gemm_case. Captures nvte_grouped_gemm on a
+// non-default stream and replays it twice, verifying outputs against the multi-tensor
+// reference after each replay. Asserts capture succeeds, replay is correct, and the
+// operation is deterministic across replays.
+void run_grouped_gemm_graph_case(const TestParams& params) {
+ if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
+ GTEST_SKIP() << reason;
+ }
+ auto ref = make_grouped_gemm_ref(params);
+ const auto& shapes = ref.shapes;
+ const size_t num_gemms = ref.num_gemms;
+
+ std::vector A_views, B_views;
+ A_views.reserve(num_gemms);
+ B_views.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ A_views.push_back(&ref.A_tensors[i]);
+ B_views.push_back(&ref.B_tensors[i]);
+ }
+ GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode());
+ GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());
+
+ std::vector C_tensors, D_group_tensors;
+ C_tensors.reserve(num_gemms);
+ D_group_tensors.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ const auto [M, N, K] = shapes[i];
+ (void)K;
+ if (!params.use_null_c) {
+ C_tensors.emplace_back(
+ Tensor("C" + std::to_string(i), std::vector{M, N}, params.output_dtype));
+ }
+ D_group_tensors.emplace_back(
+ Tensor("D_group" + std::to_string(i), std::vector{M, N}, params.output_dtype));
+ }
+
+ std::vector C_views, D_views;
+ for (size_t i = 0; i < num_gemms; ++i) {
+ if (!params.use_null_c) C_views.push_back(&C_tensors[i]);
+ D_views.push_back(&D_group_tensors[i]);
+ }
+
+ std::optional grouped_C;
+ if (!params.use_null_c) {
+ grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING);
+ }
+ GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING);
+
+ AlphaBetaTensors ab = make_alpha_beta(num_gemms);
+ const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
+ Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte);
+ Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte);
+ GroupedMatmulConfigWrapper grouped_config;
+ if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true);
+
+ cudaStream_t stream;
+ NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+ auto body = [&](cudaStream_t s) {
+ for (auto& d : D_group_tensors) {
+ NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0,
+ bytes(d.rowwise_shape(), d.dtype()), s));
+ }
+ nvte_grouped_gemm(grouped_A.get_handle(), params.transa, grouped_B.get_handle(),
+ params.transb,
+ params.use_null_c ? nullptr : grouped_C->get_handle(),
+ grouped_D.get_handle(), ab.alpha.data(), ab.beta.data(),
+ setup_ws.data(), cublas_ws.data(), grouped_config, s);
+ };
+
+ capture_and_replay_grouped_gemm(
+ stream, /*n_replays=*/2, body,
+ [&](int iter) {
+ const std::string tag = "grouped_graph_replay_" + std::to_string(iter);
+ compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, tag.c_str());
+ });
+
+ NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
+}
+
+// Graph-capture variant of run_grouped_gemm_discrete_out_case.
+void run_grouped_gemm_discrete_out_graph_case(const TestParams& params) {
+ if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
+ GTEST_SKIP() << reason;
+ }
+ auto ref = make_grouped_gemm_ref(params);
+ const auto& shapes = ref.shapes;
+ const size_t num_gemms = ref.num_gemms;
+
+ std::vector A_views, B_views;
+ A_views.reserve(num_gemms);
+ B_views.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ A_views.push_back(&ref.A_tensors[i]);
+ B_views.push_back(&ref.B_tensors[i]);
+ }
+ GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode());
+ GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());
+
+ std::vector C_tensors, D_list_tensors;
+ C_tensors.reserve(num_gemms);
+ D_list_tensors.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ const auto [M, N, K] = shapes[i];
+ (void)K;
+ if (!params.use_null_c) {
+ C_tensors.emplace_back(
+ Tensor("C" + std::to_string(i), std::vector{M, N}, params.output_dtype));
+ }
+ D_list_tensors.emplace_back(
+ Tensor("D_list" + std::to_string(i), std::vector{M, N}, params.output_dtype));
+ }
+
+ std::vector C_list_ptrs, D_list_ptrs;
+ if (!params.use_null_c) C_list_ptrs.reserve(num_gemms);
+ D_list_ptrs.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ if (!params.use_null_c) C_list_ptrs.push_back(C_tensors[i].data());
+ D_list_ptrs.push_back(D_list_tensors[i].data());
+ }
+
+ AlphaBetaTensors ab = make_alpha_beta(num_gemms);
+ const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
+ Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte);
+ Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte);
+ GroupedMatmulConfigWrapper grouped_config;
+ if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true);
+
+ cudaStream_t stream;
+ NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+ auto body = [&](cudaStream_t s) {
+ for (auto& d : D_list_tensors) {
+ NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0,
+ bytes(d.rowwise_shape(), d.dtype()), s));
+ }
+ nvte_grouped_gemm_with_discrete_out(
+ grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb,
+ params.use_null_c ? nullptr : C_list_ptrs.data(),
+ params.use_null_c ? 0 : num_gemms, D_list_ptrs.data(), num_gemms, ab.alpha.data(),
+ ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, s);
+ };
+
+ auto verify = [&](int iter) {
+ const std::string tag = "discrete_out_graph_replay_" + std::to_string(iter);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ D_list_tensors[i].to_cpu();
+ ref.D_multi[i].to_cpu();
+ auto [atol, rtol] = getTolerances(ref.D_multi[i].dtype());
+ switch (ref.D_multi[i].dtype()) {
+ case DType::kBFloat16:
+ compareResults(tag.c_str(), D_list_tensors[i],
+ ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol);
+ break;
+ case DType::kFloat16:
+ compareResults(tag.c_str(), D_list_tensors[i],
+ ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol);
+ break;
+ case DType::kFloat32:
+ compareResults(tag.c_str(), D_list_tensors[i],
+ ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol);
+ break;
+ default:
+ NVTE_ERROR("Unsupported D dtype in test: " +
+ std::to_string(static_cast(ref.D_multi[i].dtype())));
+ }
+ }
+ };
+
+ capture_and_replay_grouped_gemm(stream, /*n_replays=*/2, body, verify);
+ NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
+}
+
+// Graph-capture variant of run_grouped_gemm_discrete_in_case.
+void run_grouped_gemm_discrete_in_graph_case(const TestParams& params) {
+ if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) {
+ GTEST_SKIP() << reason;
+ }
+ auto ref = make_grouped_gemm_ref(params);
+ const auto& shapes = ref.shapes;
+ const size_t num_gemms = ref.num_gemms;
+
+ std::vector B_views;
+ B_views.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) B_views.push_back(&ref.B_tensors[i]);
+ GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode());
+
+ std::vector C_tensors, D_group_tensors;
+ C_tensors.reserve(num_gemms);
+ D_group_tensors.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) {
+ const auto [M, N, K] = shapes[i];
+ (void)K;
+ if (!params.use_null_c) {
+ C_tensors.emplace_back(Tensor("C" + std::to_string(i),
+ std::vector{M, N}, params.output_dtype));
+ }
+ D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i),
+ std::vector{M, N}, params.output_dtype));
+ }
+
+ std::vector C_views, D_views;
+ for (size_t i = 0; i < num_gemms; ++i) {
+ if (!params.use_null_c) C_views.push_back(&C_tensors[i]);
+ D_views.push_back(&D_group_tensors[i]);
+ }
+
+ std::optional grouped_C;
+ if (!params.use_null_c) {
+ grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING);
+ }
+ GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING);
+
+ AlphaBetaTensors ab = make_alpha_beta(num_gemms);
+ const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms);
+ Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte);
+ Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte);
+
+ std::vector A_list_ptrs;
+ A_list_ptrs.reserve(num_gemms);
+ for (size_t i = 0; i < num_gemms; ++i) A_list_ptrs.push_back(ref.A_tensors[i].data());
+
+ GroupedMatmulConfigWrapper grouped_config;
+ if (ref.use_split_accum) grouped_config.set_use_split_accumulator(true);
+
+ cudaStream_t stream;
+ NVTE_CHECK_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+
+ auto body = [&](cudaStream_t s) {
+ for (auto& d : D_group_tensors) {
+ NVTE_CHECK_CUDA(cudaMemsetAsync(d.rowwise_dptr(), 0,
+ bytes(d.rowwise_shape(), d.dtype()), s));
+ }
+ nvte_grouped_gemm_with_discrete_inputA(
+ A_list_ptrs.data(), num_gemms, params.transa, grouped_B.get_handle(), params.transb,
+ params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(),
+ ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, s);
+ };
+
+ capture_and_replay_grouped_gemm(
+ stream, /*n_replays=*/2, body,
+ [&](int iter) {
+ const std::string tag = "discrete_in_graph_replay_" + std::to_string(iter);
+ compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, tag.c_str());
+ });
+
+ NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
+}
+
class GroupedGemmTest : public ::testing::TestWithParam {};
TEST_P(GroupedGemmTest, CompareWithMultiTensorGemm) {
@@ -663,6 +942,18 @@ TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteIn) {
run_grouped_gemm_discrete_in_case(GetParam());
}
+TEST_P(GroupedGemmTest, CudaGraphCapture) {
+ run_grouped_gemm_graph_case(GetParam());
+}
+
+TEST_P(GroupedGemmTest, CudaGraphCaptureDiscreteOut) {
+ run_grouped_gemm_discrete_out_graph_case(GetParam());
+}
+
+TEST_P(GroupedGemmTest, CudaGraphCaptureDiscreteIn) {
+ run_grouped_gemm_discrete_in_graph_case(GetParam());
+}
+
std::string MakeGroupedGemmTestName(const testing::TestParamInfo& info) {
constexpr const char* kShapeNames[] = {"AllSameMul128", "SameMMul128", "SameNMul128",
"AllDiffMul128", "AllSameMul32"};
diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh
index d486d45f8a..da293dadfd 100755
--- a/tests/cpp_distributed/run_test_ep.sh
+++ b/tests/cpp_distributed/run_test_ep.sh
@@ -35,6 +35,12 @@ if (( MIN_SM > 0 && MIN_SM < 90 )); then
exit 0
fi
+# NCCL EP requires active NVLink P2P among ranks on the node.
+if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then
+ echo "NVLink not detected on this platform; SKIPPING."
+ exit 0
+fi
+
TEST_BIN="${BUILD_DIR}/test_ep"
if [[ ! -x "${TEST_BIN}" ]]; then
echo "ERROR: binary not found: ${TEST_BIN}"
diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu
index c7fee7720c..7dbbcdce9d 100644
--- a/tests/cpp_distributed/test_ep.cu
+++ b/tests/cpp_distributed/test_ep.cu
@@ -60,7 +60,7 @@ static std::vector generate_tokens(int rank, int num_tokens, int hidden_dim)
return v;
}
-static std::vector expected_token_counts(
+static std::vector expected_recv_tokens_per_expert(
int recv_rank, int num_processes, int num_tokens, int top_k,
int num_experts, int num_local_experts) {
int base = recv_rank * num_local_experts;
@@ -128,7 +128,7 @@ struct EPBuffers {
DevBuf topk_idx;
DevBuf topk_weights;
DevBuf tokens;
- DevBuf token_counts;
+ DevBuf recv_tokens_per_expert;
DevBuf handle_mem;
DevBuf recv_tokens;
DevBuf recv_topk_weights;
@@ -144,22 +144,26 @@ struct EPBuffers {
size_t recv_capacity = 0;
int top_k_ = 0;
size_t alignment_ = 0;
+ NVTEEpLayerConfig layer_cfg_{};
void alloc(int num_tokens, int top_k, int hidden_dim, int num_local_experts,
int ep_size, int max_tokens_per_rank, size_t alignment = 0) {
top_k_ = top_k;
alignment_ = alignment;
+ layer_cfg_ = NVTE_EP_LAYER_CONFIG_INIT;
+ layer_cfg_.top_k = top_k;
+ layer_cfg_.dispatch_output_per_expert_alignment = alignment;
recv_capacity = static_cast(ep_size) * max_tokens_per_rank * 2;
topk_idx.alloc(num_tokens * top_k);
topk_weights.alloc(num_tokens * top_k);
tokens.alloc(num_tokens * hidden_dim);
- token_counts.alloc(num_local_experts);
+ recv_tokens_per_expert.alloc(num_local_experts);
recv_tokens.alloc(recv_capacity * hidden_dim);
recv_topk_weights.alloc(recv_capacity);
result.alloc(num_tokens * hidden_dim);
- handle_mem_size = nvte_ep_handle_mem_size(NVTEEpLayerConfig{top_k, alignment});
+ handle_mem_size = nvte_ep_handle_mem_size(&layer_cfg_);
handle_mem.alloc(handle_mem_size);
grad_result.alloc(num_tokens * hidden_dim);
@@ -174,25 +178,29 @@ struct EPBuffers {
// expects.
template
struct EPTensors {
- TensorWrapper topk_idx, topk_weights, token_counts, handle_mem, tokens;
+ TensorWrapper topk_idx, topk_weights, recv_tokens_per_expert, handle_mem, tokens;
TensorWrapper recv_tokens, recv_topk_weights, result;
TensorWrapper grad_result, grad_expert, grad_tokens;
TensorWrapper g_recv_topk_weights, grad_topk_weights;
int top_k_ = 0;
size_t alignment_ = 0;
+ NVTEEpLayerConfig layer_cfg_{};
EPTensors(EPBuffers& b, int num_tokens, int top_k, int hidden_dim,
int num_local_experts) {
top_k_ = top_k;
alignment_ = b.alignment_;
+ layer_cfg_ = NVTE_EP_LAYER_CONFIG_INIT;
+ layer_cfg_.top_k = top_k;
+ layer_cfg_.dispatch_output_per_expert_alignment = b.alignment_;
constexpr DType kTokDType = test::TypeInfo::dtype;
using Shape = std::vector;
topk_idx = TensorWrapper(b.topk_idx.get(),
Shape{(size_t)num_tokens, (size_t)top_k}, DType::kInt64);
topk_weights = TensorWrapper(b.topk_weights.get(),
Shape{(size_t)num_tokens, (size_t)top_k}, DType::kFloat32);
- token_counts = TensorWrapper(b.token_counts.get(),
+ recv_tokens_per_expert = TensorWrapper(b.recv_tokens_per_expert.get(),
Shape{(size_t)num_local_experts}, DType::kInt32);
handle_mem = TensorWrapper(b.handle_mem.get(),
Shape{b.handle_mem_size}, DType::kByte);
@@ -259,7 +267,7 @@ class EpOpTestBase : public ::testing::Test {
template
int read_total_recv(const EPBuffers& buf) const {
std::vector cnt(num_local_experts_);
- NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(),
+ NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.recv_tokens_per_expert.get(),
num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost));
int total = 0;
for (int c : cnt) total += c;
@@ -300,7 +308,7 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) {
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(),
NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{},
@@ -309,9 +317,9 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) {
// 1. Per-expert counts.
std::vector got_counts(num_local_experts_);
- NVTE_CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.token_counts.get(),
+ NVTE_CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.recv_tokens_per_expert.get(),
num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost));
- auto exp_counts = expected_token_counts(g_process_id, g_num_processes, num_tokens_, top_k_,
+ auto exp_counts = expected_recv_tokens_per_expert(g_process_id, g_num_processes, num_tokens_, top_k_,
num_experts_, num_local_experts_);
int total_recv = 0;
for (int i = 0; i < num_local_experts_; ++i) {
@@ -379,7 +387,7 @@ TYPED_TEST(EPCombineTest, Combine) {
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(),
NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{},
@@ -426,7 +434,7 @@ TYPED_TEST(EPCombineBwdTest, CombineBwdCheck) {
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(),
NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{},
@@ -447,7 +455,7 @@ TYPED_TEST(EPCombineBwdTest, CombineBwdCheck) {
int total_recv = this->template read_total_recv(buf);
std::vector cnt(num_local_experts_);
- NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(),
+ NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.recv_tokens_per_expert.get(),
num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost));
std::vector h_ge(buf.recv_capacity * hidden_dim_);
NVTE_CHECK_CUDA(cudaMemcpy(h_ge.data(), buf.grad_expert.get(),
@@ -495,7 +503,7 @@ TYPED_TEST(EPDispatchBwdTest, DispatchBwdCheck) {
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(),
NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{},
@@ -563,7 +571,7 @@ TYPED_TEST(EPDispatchBwdGradWeightsTest, RoundTrip) {
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
NVTE_CHECK_CUDA(cudaMemsetAsync(buf.recv_topk_weights.get(), 0,
buf.recv_topk_weights.bytes(), stream));
ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
@@ -634,7 +642,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(),
NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{},
@@ -759,7 +767,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) {
cudaStream_t stream;
NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
- ASSERT_NO_THROW(nvte_ep_prepare(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.token_counts.data(), NVTEEpLayerConfig{ref_t.top_k_, ref_t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.recv_tokens_per_expert.data(), nullptr, &ref_t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(ref_t.handle_mem.data(), ref_t.topk_idx.data(),
ref_t.tokens.data(), NVTECommWindow{}, ref_t.topk_weights.data(),
NVTECommWindow{}, ref_t.recv_tokens.data(), NVTECommWindow{},
@@ -800,7 +808,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) {
sym_t.recv_tokens = TensorWrapper(sym_recv.ptr,
std::vector{sym_buf.recv_capacity, (size_t)hidden_dim_}, kTokDType);
- ASSERT_NO_THROW(nvte_ep_prepare(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.token_counts.data(), NVTEEpLayerConfig{sym_t.top_k_, sym_t.alignment_}, stream));
+ ASSERT_NO_THROW(nvte_ep_prepare(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.recv_tokens_per_expert.data(), nullptr, &sym_t.layer_cfg_, stream));
ASSERT_NO_THROW(nvte_ep_dispatch(sym_t.handle_mem.data(), sym_t.topk_idx.data(),
sym_t.tokens.data(), symm_window(sym_tokens),
sym_t.topk_weights.data(), NVTECommWindow{},
diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h
index d5e006cef6..7cf6017090 100644
--- a/tests/cpp_distributed/test_ep_common.h
+++ b/tests/cpp_distributed/test_ep_common.h
@@ -146,7 +146,7 @@ static bool ep_bootstrap(int argc, char* argv[]) {
ncclUniqueId uid{};
exchange_unique_id(&uid);
- NVTEEpGroupConfig group_config{};
+ NVTEEpGroupConfig group_config = NVTE_EP_GROUP_CONFIG_INIT;
group_config.ep_size = g_ep_size;
group_config.num_experts = g_num_experts;
group_config.max_tokens_per_rank = g_max_tokens_per_rank;
@@ -156,7 +156,7 @@ static bool ep_bootstrap(int argc, char* argv[]) {
group_config.max_token_dtype = g_max_token_dtype;
NVTE_CHECK_NCCL(ncclCommInitRank(&g_ep_comm, g_num_processes, uid, g_process_id));
- nvte_ep_initialize(static_cast(g_ep_comm), group_config);
+ nvte_ep_initialize(static_cast(g_ep_comm), &group_config);
if (g_process_id == 0) {
printf("EP initialized: ep_size=%d num_experts=%d "
@@ -173,7 +173,7 @@ static bool ep_bootstrap(int argc, char* argv[]) {
static void ep_reinitialize(int zero_copy) {
if (!g_ep_initialized) return;
nvte_ep_shutdown();
- NVTEEpGroupConfig group_config{};
+ NVTEEpGroupConfig group_config = NVTE_EP_GROUP_CONFIG_INIT;
group_config.ep_size = g_ep_size;
group_config.num_experts = g_num_experts;
group_config.max_tokens_per_rank = g_max_tokens_per_rank;
@@ -181,7 +181,7 @@ static void ep_reinitialize(int zero_copy) {
group_config.hidden_dim = g_hidden_dim;
group_config.max_token_dtype = g_max_token_dtype;
group_config.zero_copy = zero_copy;
- nvte_ep_initialize(static_cast(g_ep_comm), group_config);
+ nvte_ep_initialize(static_cast(g_ep_comm), &group_config);
}
// Tear down in dependency order: backend's ep_group reads from ep_comm,
diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh
index d32ce5f5d3..ff89f712eb 100755
--- a/tests/jax/multi_process_launch_ep.sh
+++ b/tests/jax/multi_process_launch_ep.sh
@@ -32,6 +32,13 @@ if [ "${NUM_RUNS}" -lt 4 ]; then
echo "NCCL EP requires at least 4 GPUs (found ${NUM_RUNS}); SKIPPING."
exit 0
fi
+
+# NCCL EP requires active NVLink P2P among ranks on the node.
+if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then
+ echo "NVLink not detected on this platform — EP test requires NVLink; SKIPPING."
+ exit 0
+fi
+
# Default test mesh is (2, 2); use exactly 4 ranks even on larger boxes.
NUM_RUNS="${NVTE_TEST_EP_NUM_RANKS:-4}"
diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py
index 3d2f99b51b..82b9df262f 100644
--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -411,6 +411,13 @@ def run_dpa_with_cp(
cu_seqlens_kv=cu_seqlens_kv,
cu_seqlens_q_padded=cu_seqlens_q_padded,
cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+ # Test runner sets cu_seqlens_q == cu_seqlens_q_padded for the
+ # FlashAttention path, i.e. no inter-sequence padding. Declare this
+ # explicitly so the sync-free auto-detect (which conservatively
+ # picks True when padded cu_seqlens are present) does not disable FA.
+ pad_between_seqs=(
+ (kernel_backend != "FlashAttention") if qkv_format == "thd" else None
+ ),
fp8_output=fp8_mha,
)
if config.return_max_logit:
@@ -528,6 +535,12 @@ def run_dpa_with_cp(
cu_seqlens_kv=cu_seqlens_kv,
cu_seqlens_q_padded=cu_seqlens_q_padded,
cu_seqlens_kv_padded=cu_seqlens_kv_padded,
+ # See note above (non-CP branch): same explicit declaration so
+ # FlashAttention isn't disabled by the conservative sync-free
+ # auto-detect when this test path constructs no inter-seq padding.
+ pad_between_seqs=(
+ (kernel_backend != "FlashAttention") if qkv_format == "thd" else None
+ ),
fp8_output=fp8_mha,
)
if config.return_max_logit:
diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py
new file mode 100644
index 0000000000..0acc00cd57
--- /dev/null
+++ b/tests/pytorch/distributed/run_ep.py
@@ -0,0 +1,521 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Multi-process PyTorch EP tests, launched via torchrun (one process per GPU)."""
+
+import os
+import sys
+import unittest
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from transformer_engine.pytorch.ep import (
+ EpBuffer,
+ ep_bootstrap,
+ ep_finalize,
+ ep_prepare,
+ ep_dispatch,
+ ep_combine,
+ symm_mem_alloc,
+ _ep_combine_raw,
+ _ep_dispatch_raw,
+)
+
+
+ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1"
+
+# Must come after the transformer_engine import so libtransformer_engine.so is loaded.
+import transformer_engine_torch as tex # noqa: F401
+
+
+NUM_LOCAL_EXPERTS = 2
+HIDDEN_DIM = 32
+TOP_K = 2
+TOKENS_PER_RANK = 4
+
+
+def _zero_copy_test_include(fn):
+ """Mark a test to also run in the zero-copy pass; others skip there."""
+ fn._zero_copy_test_include = True
+ return fn
+
+
+class _StageToSymm(torch.autograd.Function):
+ """Identity op that stages ``src`` into a symm-mem buffer; grad passes through.
+ Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine.
+ """
+
+ @staticmethod
+ def forward(ctx, src, symm_buf): # type: ignore[override]
+ symm_buf.copy_(src)
+ return symm_buf
+
+ @staticmethod
+ def backward(ctx, g): # type: ignore[override]
+ return g, None
+
+
+class _GradToSymm(torch.autograd.Function):
+ """Identity fwd; bwd stages the upstream grad into a symm-mem buffer and
+ returns it, so the next backward (dispatch_bwd) receives a symm-window grad
+ input — which zero-copy ncclEpCombine requires.
+ """
+
+ @staticmethod
+ def forward(ctx, x, symm_buf): # type: ignore[override]
+ ctx.symm_buf = symm_buf
+ return x
+
+ @staticmethod
+ def backward(ctx, g): # type: ignore[override]
+ ctx.symm_buf.copy_(g)
+ return ctx.symm_buf, None
+
+
+def _device_sm() -> int:
+ major, minor = torch.cuda.get_device_capability()
+ return major * 10 + minor
+
+
+def _build_ep_group():
+ """EP group spanning all ranks of the default PG."""
+ world_pg = dist.distributed_c10d._get_default_group()
+ ranks = list(range(world_pg.size()))
+ return dist.new_group(ranks=ranks, backend="nccl")
+
+
+def _make_identity_inputs(rank, ep_size, device="cuda"):
+ """Per-rank identity routing + uniform weights so combine matches tokens."""
+ T = TOKENS_PER_RANK
+ E = ep_size * NUM_LOCAL_EXPERTS
+ topk_idx = np.empty((T, TOP_K), dtype=np.int64)
+ base = rank * T
+ for t in range(T):
+ for k in range(TOP_K):
+ topk_idx[t, k] = ((base + t) * TOP_K + k) % E
+ tokens_np = np.linspace(
+ 0.1 + rank * 0.01, 0.9 + rank * 0.01, T * HIDDEN_DIM, dtype=np.float32
+ ).reshape(T, HIDDEN_DIM)
+ topk_weights = np.full((T, TOP_K), 1.0 / TOP_K, dtype=np.float32)
+ return (
+ torch.from_numpy(topk_idx).to(device),
+ torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16),
+ torch.from_numpy(topk_weights).to(device),
+ )
+
+
+class _Cfg:
+ rank: int
+ world_size: int
+ ep_size: int
+ num_experts: int
+ recv_capacity_per_rank: int
+ device: torch.device
+
+
+def _make_cfg() -> _Cfg:
+ cfg = _Cfg()
+ cfg.rank = dist.get_rank()
+ cfg.world_size = dist.get_world_size()
+ cfg.ep_size = cfg.world_size
+ cfg.num_experts = NUM_LOCAL_EXPERTS * cfg.ep_size
+ T = TOKENS_PER_RANK
+ active = min(cfg.num_experts, T * cfg.ep_size * TOP_K)
+ overconc = cfg.num_experts // active
+ cfg.recv_capacity_per_rank = NUM_LOCAL_EXPERTS * max(T * cfg.ep_size * TOP_K, 16) * overconc * 2
+ cfg.device = torch.device("cuda", torch.cuda.current_device())
+ return cfg
+
+
+class TestEP(unittest.TestCase):
+ cfg: _Cfg
+ ep_group: dist.ProcessGroup
+
+ @classmethod
+ def setUpClass(cls):
+ if _device_sm() < 90:
+ raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})")
+ cls.cfg = _make_cfg()
+ cls.ep_group = _build_ep_group()
+ ep_bootstrap(
+ cls.ep_group,
+ num_experts=cls.cfg.num_experts,
+ max_tokens_per_rank=TOKENS_PER_RANK,
+ recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank,
+ hidden_dim=HIDDEN_DIM,
+ zero_copy=ZERO_COPY,
+ )
+
+ def setUp(self):
+ # Only the zero-copy-capable tests run in the zero-copy pass.
+ if ZERO_COPY and not getattr(
+ getattr(self, self._testMethodName), "_zero_copy_test_include", False
+ ):
+ self.skipTest("not exercised in zero-copy mode")
+
+ def _make_buffer(
+ self,
+ alignment=0,
+ top_k=TOP_K,
+ dispatch_recv_tokens=None,
+ combine_grad_expert_out=None,
+ ):
+ return EpBuffer(
+ top_k=top_k,
+ max_tokens_per_rank=TOKENS_PER_RANK,
+ recv_capacity_per_rank=self.cfg.recv_capacity_per_rank,
+ hidden_dim=HIDDEN_DIM,
+ num_local_experts=NUM_LOCAL_EXPERTS,
+ alignment=alignment,
+ dispatch_recv_tokens=dispatch_recv_tokens,
+ combine_grad_expert_out=combine_grad_expert_out,
+ )
+
+ def _expert_out(self, expert_out):
+ """Stage the combine input into symm-mem under zero-copy (combine requires it)."""
+ if not ZERO_COPY:
+ return expert_out
+ symm_buf = symm_mem_alloc(tuple(expert_out.shape), expert_out.dtype, self.ep_group)
+ return _StageToSymm.apply(expert_out, symm_buf)
+
+ def _stage_grad_symm(self, x, symm_buf=None):
+ """Route x's upstream grad through a symm-mem buffer so dispatch_bwd gets
+ a symm-window grad input under zero-copy; passthrough otherwise. Pass a
+ pre-allocated symm_buf to avoid allocating during an interleaved schedule."""
+ if not ZERO_COPY:
+ return x
+ if symm_buf is None:
+ symm_buf = symm_mem_alloc(tuple(x.shape), x.dtype, self.ep_group)
+ return _GradToSymm.apply(x, symm_buf)
+
+ def _make_raw_recv(self, dtype=torch.bfloat16):
+ """Raw recv tensors + token_counts for the primitive tests."""
+ rc = self.cfg.recv_capacity_per_rank
+ return (
+ torch.empty(rc, HIDDEN_DIM, dtype=dtype, device=self.cfg.device),
+ torch.empty(rc, dtype=torch.float32, device=self.cfg.device),
+ torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int32, device=self.cfg.device),
+ )
+
+ @staticmethod
+ def _weighted(recv_tokens, recv_w):
+ """fp32 per-slot weighting + cast back; matches the upstream combine input."""
+ mask = (recv_w != 0).to(torch.float32).unsqueeze(-1)
+ return (recv_tokens.float() * recv_w.unsqueeze(-1).float() * mask).to(recv_tokens.dtype)
+
+ def _moe_step(self, buffer, topk_idx, tokens, w):
+ recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, w)
+ expert_out = self._weighted(recv_t, recv_w_out)
+ return ep_combine(buffer, expert_out)
+
+ # Prepare
+
+ def test_primitive_prepare(self):
+ buf = self._make_buffer()
+ topk_idx, _toks, _w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ token_counts = ep_prepare(buf, topk_idx)
+ torch.cuda.synchronize()
+ self.assertEqual(token_counts.shape, (NUM_LOCAL_EXPERTS,))
+ local = int(token_counts.sum().item())
+ total = torch.tensor([local], dtype=torch.int64, device=self.cfg.device)
+ dist.all_reduce(total, op=dist.ReduceOp.SUM, group=self.ep_group)
+ self.assertEqual(int(total.item()), self.cfg.world_size * TOKENS_PER_RANK * TOP_K)
+
+ # Identity round-trip via raw primitives
+
+ def test_primitive_dispatch_combine_identity(self):
+ buf = self._make_buffer()
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ recv_tokens, recv_w, _ = self._make_raw_recv()
+ ep_prepare(buf, topk_idx)
+ _ep_dispatch_raw(buf, topk_idx, tokens, w, recv_tokens, recv_w)
+ result = torch.empty_like(tokens)
+ _ep_combine_raw(buf, self._weighted(recv_tokens, recv_w), result)
+ torch.cuda.synchronize()
+ torch.testing.assert_close(result.float(), tokens.float(), atol=5e-2, rtol=5e-2)
+
+ # Autograd
+
+ @_zero_copy_test_include
+ def test_dispatch_autograd(self):
+ """0.5*||recv_tokens||^2 ; grad_tokens equals TOP_K * tokens. Covers the
+ EpBuffer-owned recv tokens (symm-mem under zero-copy) and, in normal
+ mode, a caller-supplied recv_tokens buffer."""
+ if ZERO_COPY:
+ cases = [("buffer_owned", None)]
+ else:
+ rt_buf, _rw_buf, _ = self._make_raw_recv()
+ cases = [
+ ("default_alloc", None),
+ ("caller_recv", rt_buf),
+ ]
+ for label, recv_tokens in cases:
+ with self.subTest(case=label):
+ buf = self._make_buffer(dispatch_recv_tokens=recv_tokens)
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ tokens_p = tokens.detach().clone().requires_grad_(True)
+ rt, rw, _tc = ep_dispatch(buf, tokens_p, topk_idx, w)
+ if recv_tokens is not None: # caller-supplied recv_tokens must be used in place
+ self.assertEqual(rt.data_ptr(), recv_tokens.data_ptr())
+ rt = self._stage_grad_symm(rt)
+ rw = self._stage_grad_symm(rw)
+ (0.5 * (rt.float() ** 2).sum() + 0.0 * rw.float().sum()).backward()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(
+ tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2
+ )
+
+ @_zero_copy_test_include
+ def test_caller_provides_dispatch_recv_tokens(self):
+ """Caller-supplied recv_tokens: EpBuffer adopts it (recv_topk_weights stays
+ owned) and ep_dispatch returns a view of the caller's buffer."""
+ if ZERO_COPY:
+ rc = self.cfg.recv_capacity_per_rank
+ rt_buf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group)
+ else:
+ rt_buf, _rw_buf, _ = self._make_raw_recv()
+ buf = self._make_buffer(dispatch_recv_tokens=rt_buf)
+ self.assertEqual(buf.recv_tokens_symm_buf.data_ptr(), rt_buf.data_ptr())
+ if ZERO_COPY: # recv_topk_weights is always buffer-owned in zero-copy
+ self.assertIsNotNone(buf.recv_topk_weights_symm_buf)
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ tokens_p = tokens.detach().clone().requires_grad_(True)
+ rt, rw, _ = ep_dispatch(buf, tokens_p, topk_idx, w)
+ self.assertEqual(rt.data_ptr(), rt_buf.data_ptr())
+ rt = self._stage_grad_symm(rt)
+ rw = self._stage_grad_symm(rw)
+ (0.5 * (rt.float() ** 2).sum() + 0.0 * rw.float().sum()).backward()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(
+ tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2
+ )
+
+ @_zero_copy_test_include
+ def test_caller_provides_grad_expert_out(self):
+ """Caller-supplied grad_expert_out: EpBuffer adopts it as the combine
+ backward grad target (symm-mem under zero-copy)."""
+ rc = self.cfg.recv_capacity_per_rank
+ if ZERO_COPY:
+ gbuf = symm_mem_alloc((rc, HIDDEN_DIM), torch.bfloat16, self.ep_group)
+ else:
+ gbuf = torch.empty(rc, HIDDEN_DIM, dtype=torch.bfloat16, device=self.cfg.device)
+ buf = self._make_buffer(combine_grad_expert_out=gbuf)
+ self.assertEqual(buf.grad_expert_out_symm_buf.data_ptr(), gbuf.data_ptr())
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ tokens_p = tokens.detach().clone().requires_grad_(True)
+ recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w)
+ recv_t = self._stage_grad_symm(recv_t)
+ recv_w = self._stage_grad_symm(recv_w)
+ expert_out = self._expert_out(self._weighted(recv_t, recv_w))
+ out = ep_combine(buf, expert_out)
+ (0.5 * (out.float() ** 2).sum()).backward()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2)
+ torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2)
+
+ # Multi-iter stability
+
+ def test_dispatch_autograd_multiple_iterations(self):
+ """5 fwd+bwd iters on the same EpBuffer must be bit-stable."""
+ buf = self._make_buffer()
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+
+ def one_step():
+ tokens_p = tokens.detach().clone().requires_grad_(True)
+ out = self._moe_step(buf, topk_idx, tokens_p, w)
+ loss = 0.5 * (out.float() ** 2).sum()
+ loss.backward()
+ return out.detach().clone(), tokens_p.grad.detach().clone()
+
+ out_ref, grad_ref = one_step()
+ torch.cuda.synchronize()
+ for _ in range(4):
+ out_i, grad_i = one_step()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(out_i, out_ref, atol=0, rtol=0)
+ torch.testing.assert_close(grad_i, grad_ref, atol=0, rtol=0)
+
+ # CUDA graph
+
+ def test_cuda_graph_capture(self):
+ """Capture raw dispatch+combine into a CUDA graph; replay must be bit-stable."""
+ buf = self._make_buffer()
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ recv_tokens, recv_w, _ = self._make_raw_recv()
+ result = torch.empty_like(tokens)
+
+ def step():
+ ep_prepare(buf, topk_idx)
+ _ep_dispatch_raw(buf, topk_idx, tokens, w, recv_tokens, recv_w)
+ _ep_combine_raw(buf, self._weighted(recv_tokens, recv_w), result)
+
+ for _ in range(3):
+ step()
+ torch.cuda.synchronize()
+
+ # Routing is fixed per layer; prepare runs once before capture.
+ ep_prepare(buf, topk_idx)
+ torch.cuda.synchronize()
+
+ graph = torch.cuda.CUDAGraph()
+ s = torch.cuda.Stream()
+ s.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(s):
+ with torch.cuda.graph(graph):
+ _ep_dispatch_raw(buf, topk_idx, tokens, w, recv_tokens, recv_w)
+ _ep_combine_raw(buf, self._weighted(recv_tokens, recv_w), result)
+ torch.cuda.current_stream().wait_stream(s)
+ torch.cuda.synchronize()
+
+ ref = result.clone()
+ for _ in range(5):
+ graph.replay()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(result.float(), ref.float(), atol=0, rtol=0)
+
+ # PP-1F1B handle isolation
+
+ @_zero_copy_test_include
+ def test_pp_1f1b_two_handles(self):
+ """PP-1F1B interleave (F0 F1 B0 F2 B1 B2) over 3 per-microbatch buffers,
+ run eagerly and replayed from a CUDA graph capturing the full fwd+bwd
+ schedule (prepare included; routing is fixed so replay reproduces it)."""
+ for capture in (False, True):
+ with self.subTest(capture=capture):
+ self._run_1f1b(capture)
+
+ def _run_1f1b(self, capture):
+ T, H = TOKENS_PER_RANK, HIDDEN_DIM
+ idx, _toks, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ scales = (0.13, 0.41, 0.77)
+ buffers, tokens, tokens_p = [], [], []
+ for s in scales:
+ buffers.append(self._make_buffer())
+ t = torch.full(
+ (T, H), s + self.cfg.rank * 0.01, dtype=torch.bfloat16, device=self.cfg.device
+ )
+ tokens.append(t)
+ tokens_p.append(t.detach().clone().requires_grad_(True))
+
+ recv = [None, None, None]
+ # Per-microbatch grad-staging buffers, symm-mem under zero-copy and
+ # pre-allocated so nothing is allocated/freed mid-interleave. The recv
+ # outputs are owned by each EpBuffer (symm-mem under zero-copy).
+ recv_w = [None, None, None]
+ rc = self.cfg.recv_capacity_per_rank
+ if ZERO_COPY:
+ gbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales]
+ gbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales]
+ else:
+ gbuf_t = gbuf_w = [None, None, None]
+
+ def fwd(k):
+ rt, rw, _ = ep_dispatch(buffers[k], tokens_p[k], idx, w)
+ recv[k] = self._stage_grad_symm(rt, gbuf_t[k])
+ recv_w[k] = self._stage_grad_symm(rw, gbuf_w[k])
+
+ def bwd(k):
+ (0.5 * (recv[k].float() ** 2).sum() + 0.0 * recv_w[k].float().sum()).backward()
+ recv[k] = None
+ recv_w[k] = None
+
+ def interleave():
+ fwd(0)
+ fwd(1)
+ bwd(0)
+ fwd(2)
+ bwd(1)
+ bwd(2)
+
+ def zero_grads():
+ for tp in tokens_p:
+ if tp.grad is not None:
+ tp.grad.zero_()
+
+ if not capture:
+ interleave()
+ else:
+ # Warmup on a side stream, then capture the full schedule and replay.
+ # Grads stay pre-allocated (zeroed, not None) so backward accumulates
+ # in place during both capture and replay.
+ s = torch.cuda.Stream()
+ s.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(s):
+ for _ in range(3):
+ zero_grads()
+ interleave()
+ torch.cuda.current_stream().wait_stream(s)
+ torch.cuda.synchronize()
+
+ zero_grads()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ interleave()
+ zero_grads()
+ graph.replay()
+
+ torch.cuda.synchronize()
+ for k in range(3):
+ torch.testing.assert_close(
+ tokens_p[k].grad.float(),
+ tokens[k].float() * float(TOP_K),
+ atol=5e-2,
+ rtol=5e-2,
+ )
+
+ @_zero_copy_test_include
+ def test_combine_autograd(self):
+ """ep_combine fwd+bwd; bwd grad target is the EpBuffer symm buffer (zc) or in-flight."""
+ buf = self._make_buffer()
+ topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+ tokens_p = tokens.detach().clone().requires_grad_(True)
+ recv_t, recv_w, _ = ep_dispatch(buf, tokens_p, topk_idx, w)
+ recv_t = self._stage_grad_symm(recv_t)
+ recv_w = self._stage_grad_symm(recv_w)
+ expert_out = self._expert_out(self._weighted(recv_t, recv_w))
+ out = ep_combine(buf, expert_out)
+ (0.5 * (out.float() ** 2).sum()).backward()
+ torch.cuda.synchronize()
+ torch.testing.assert_close(out.float(), tokens.float(), atol=5e-2, rtol=5e-2)
+ torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2)
+
+ # Input validation
+
+ def test_topk_int32_raises_clear_error(self):
+ buf = self._make_buffer()
+ topk_idx_int32 = torch.zeros(
+ TOKENS_PER_RANK, TOP_K, dtype=torch.int32, device=self.cfg.device
+ )
+ with self.assertRaises(RuntimeError) as cm:
+ ep_prepare(buf, topk_idx_int32)
+ msg = str(cm.exception)
+ self.assertIn("topk_idx", msg)
+ self.assertIn(".long()", msg)
+
+
+def _init_distributed():
+ dist.init_process_group(backend="nccl")
+ torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+ try:
+ from torch.distributed import _symmetric_memory as _symm_mem
+
+ _symm_mem.set_backend("NCCL")
+ except (ImportError, RuntimeError):
+ pass
+
+
+if __name__ == "__main__":
+ _init_distributed()
+ loader = unittest.TestLoader()
+ name_filter = os.environ.get("NVTE_EP_TEST_FILTER")
+ if name_filter:
+ loader.testMethodPrefix = name_filter
+ suite = loader.loadTestsFromTestCase(TestEP)
+ runner = unittest.TextTestRunner(stream=sys.stdout, verbosity=2)
+ result = runner.run(suite)
+ dist.barrier()
+ ep_finalize()
+ dist.destroy_process_group()
+ sys.exit(0 if result.wasSuccessful() else 1)
diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh
new file mode 100755
index 0000000000..68b691f787
--- /dev/null
+++ b/tests/pytorch/distributed/run_test_ep.sh
@@ -0,0 +1,74 @@
+#!/bin/bash
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+#
+# Launcher for tests/pytorch/distributed/run_ep.py. Auto-detects GPU count.
+# Short timeout by default to surface hangs early.
+
+set -uo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
+
+DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
+if [ "${DETECTED_GPUS}" -lt 4 ]; then
+ echo "EP requires >= 4 GPUs (found ${DETECTED_GPUS}); SKIPPING."
+ exit 0
+fi
+
+# NCCL EP requires active NVLink P2P among ranks on the node.
+# On PCIe-only nodes (no NVLink) it falls back to the network
+# transport and deadlocks, so skip cleanly there.
+if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then
+ echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING."
+ exit 0
+fi
+
+NUM_RANKS="${NVTE_TEST_EP_NUM_RANKS:-${DETECTED_GPUS}}"
+if [ "${NUM_RANKS}" -gt 8 ]; then NUM_RANKS=8; fi
+
+# Short timeout to detect hangs early.
+TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-120}"
+
+# Stage NCCL EP JIT cubins on tmpfs to keep iteration fast.
+: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"}
+export NCCL_EP_JIT_CACHE_DIR
+mkdir -p "$NCCL_EP_JIT_CACHE_DIR"
+
+SCRIPT="${SCRIPT_DIR}/run_ep.py"
+
+RET=0
+
+# Run the suite once per IO mode. Modes can't be mixed in one process
+# (ep_bootstrap is once-per-process), so zero-copy gets its own run; only the
+# zero-copy-capable tests execute there (the rest self-skip).
+run_pass() {
+ local label="$1"
+ local zc="$2"
+ local log="stdout_ep_${label}.txt"
+ echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ==="
+ # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun.
+ NVTE_EP_ZERO_COPY="${zc}" setsid timeout --foreground --kill-after=10 --signal=TERM \
+ "${TEST_TIMEOUT_S}" \
+ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \
+ "${SCRIPT}" 2>&1 | tee "${log}"
+ local rc=${PIPESTATUS[0]}
+ pkill -9 -f "tests/pytorch/distributed/run_ep.py" 2>/dev/null || true
+
+ if [ "${rc}" -ne 0 ]; then echo "[${label}] torchrun exited with ${rc}"; RET=1; fi
+ # Match unittest failure markers and unhandled Python tracebacks; torchrun
+ # prefixes per-rank stderr with "[rankN]:" so don't anchor at column 0.
+ if grep -qE "(^|]:)FAILED|(^|]:)Traceback" "${log}"; then RET=1; fi
+ if ! grep -qE "Ran [0-9]+ test|^OK$" "${log}"; then
+ echo "[${label}] ERROR: no test summary — likely hang or early crash"
+ RET=1
+ fi
+ if [ -z "${KEEP_EP_LOGS:-}" ]; then rm -f "${log}"; fi
+}
+
+run_pass "default" 0
+run_pass "zero_copy" 1
+
+exit $RET
diff --git a/tests/pytorch/distributed/test_ep.py b/tests/pytorch/distributed/test_ep.py
new file mode 100644
index 0000000000..81eef9a3c1
--- /dev/null
+++ b/tests/pytorch/distributed/test_ep.py
@@ -0,0 +1,31 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Pytest driver — spawns run_ep.py under torchrun and asserts the suite passed."""
+
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import torch
+
+TEST_ROOT = Path(__file__).parent.resolve()
+WORKER = TEST_ROOT / "run_ep.py"
+LAUNCHER = TEST_ROOT / "run_test_ep.sh"
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="EP requires >= 4 GPUs")
+def test_multi_process_ep():
+ """Launch the EP unit-test suite across all visible GPUs.
+
+ Short timeout so a hang on any rank surfaces fast rather than burning CI time.
+ """
+ timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180"))
+ proc = subprocess.run(
+ ["bash", str(LAUNCHER)],
+ env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)},
+ timeout=timeout_s + 30,
+ check=False,
+ )
+ assert proc.returncode == 0, f"EP test suite failed (rc={proc.returncode})"
diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py
index 5e6f36e8b4..c0acf2e6b3 100644
--- a/tests/pytorch/test_backward_override.py
+++ b/tests/pytorch/test_backward_override.py
@@ -858,6 +858,218 @@ def test_backward_override_recipe_matches_requested_mode(
assert quant_recipe.backward_override is None
+@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
+@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias"))
+def test_linear_backward_override_dequantized_ignores_save_original_input(
+ recipe_name: str,
+ use_bias: bool,
+) -> None:
+ reset_rng_states()
+ dtype = torch.bfloat16
+ input_shape = (32, 128)
+ out_features = 128
+ _maybe_skip_recipe_dtype(recipe_name, dtype, "linear")
+ _maybe_skip_unsupported_recipe_module_combo(recipe_name, "linear")
+ _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, "linear")
+
+ mode_recipe = make_recipe(recipe_name, backward_override="dequantized")
+ skip_unsupported_backward_override("linear", mode_recipe, "dequantized")
+
+ module_ref = te.Linear(
+ input_shape[-1],
+ out_features,
+ bias=use_bias,
+ params_dtype=dtype,
+ device="cuda",
+ save_original_input=False,
+ )
+ module_test = te.Linear(
+ input_shape[-1],
+ out_features,
+ bias=use_bias,
+ params_dtype=dtype,
+ device="cuda",
+ save_original_input=True,
+ )
+ _copy_named_parameters(module_ref, module_test)
+
+ x = torch.randn(*input_shape, dtype=dtype, device="cuda")
+ dy = torch.randn(input_shape[0], out_features, dtype=dtype, device="cuda")
+
+ y_ref, dx_ref, dw_ref, db_ref = _run_single_step(module_ref, x, dy, mode_recipe)
+ y_test, x_test, saved_operands = _run_single_step_with_saved_operands(
+ module_test, x, mode_recipe
+ )
+ _assert_saved_quantized_operand_uses_rowwise_only(saved_operands[0], name="linear_input")
+
+ y_test_detached = y_test.detach().clone()
+ y_test.backward(dy)
+ assert x_test.grad is not None
+ assert module_test.weight.grad is not None
+ dx_test = x_test.grad.detach().clone()
+ dw_test = module_test.weight.grad.detach().clone()
+ test_bias = getattr(module_test, "bias", None)
+ db_test = (
+ None if test_bias is None or test_bias.grad is None else test_bias.grad.detach().clone()
+ )
+
+ assert_close(y_test_detached, y_ref, rtol=0, atol=0, check_dtype=True)
+ assert_close(dx_test, dx_ref, rtol=0, atol=0, check_dtype=True)
+ assert_close(dw_test, dw_ref, rtol=0, atol=0, check_dtype=True)
+ if use_bias:
+ assert db_test is not None and db_ref is not None
+ assert_close(db_test, db_ref, rtol=0, atol=0, check_dtype=True)
+
+
+@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
+@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias"))
+def test_grouped_linear_backward_override_dequantized_ignores_save_original_input(
+ recipe_name: str,
+ use_bias: bool,
+) -> None:
+ reset_rng_states()
+ dtype = torch.bfloat16
+ in_features = 128
+ out_features = 128
+ m_splits = [64, 64]
+ num_gemms = len(m_splits)
+ num_tokens = sum(m_splits)
+ _maybe_skip_recipe_dtype(recipe_name, dtype, "grouped_linear")
+ _maybe_skip_unsupported_recipe_module_combo(recipe_name, "grouped_linear")
+ _maybe_skip_unsupported_grouped_splits(recipe_name, m_splits)
+
+ mode_recipe = make_recipe(recipe_name, backward_override="dequantized")
+ skip_unsupported_backward_override("grouped_linear", mode_recipe, "dequantized")
+
+ module_ref = te.GroupedLinear(
+ num_gemms,
+ in_features,
+ out_features,
+ bias=use_bias,
+ params_dtype=dtype,
+ device="cuda",
+ save_original_input=False,
+ )
+ module_test = te.GroupedLinear(
+ num_gemms,
+ in_features,
+ out_features,
+ bias=use_bias,
+ params_dtype=dtype,
+ device="cuda",
+ save_original_input=True,
+ )
+ _copy_named_parameters(module_ref, module_test)
+
+ x = torch.randn(num_tokens, in_features, dtype=dtype, device="cuda")
+ dy = torch.randn(num_tokens, out_features, dtype=dtype, device="cuda")
+
+ y_ref, dx_ref, dw_ref, db_ref = _run_grouped_linear_single_step(
+ module_ref, x, m_splits, dy, mode_recipe
+ )
+ y_test, x_test, saved_operands = _run_grouped_linear_step_with_saved_operands(
+ module_test, x, m_splits, mode_recipe
+ )
+ saved_inputs = saved_operands[:num_gemms]
+ for i, saved_input in enumerate(saved_inputs):
+ _assert_saved_quantized_operand_uses_rowwise_only(
+ saved_input, name=f"grouped_linear_input{i}"
+ )
+
+ y_test_detached = y_test.detach().clone()
+ y_test.backward(dy)
+ assert x_test.grad is not None
+ dx_test = x_test.grad.detach().clone()
+ dw_test = [getattr(module_test, f"weight{i}").grad.detach().clone() for i in range(num_gemms)]
+ db_test: list[Optional[torch.Tensor]] = []
+ for i in range(num_gemms):
+ if use_bias:
+ db_test.append(getattr(module_test, f"bias{i}").grad.detach().clone())
+ else:
+ db_test.append(None)
+
+ assert_close(y_test_detached, y_ref, rtol=0, atol=0, check_dtype=True)
+ assert_close(dx_test, dx_ref, rtol=0, atol=0, check_dtype=True)
+ for test_dw, ref_dw in zip(dw_test, dw_ref):
+ assert_close(test_dw, ref_dw, rtol=0, atol=0, check_dtype=True)
+ if use_bias:
+ for test_db, ref_db in zip(db_test, db_ref):
+ assert test_db is not None and ref_db is not None
+ assert_close(test_db, ref_db, rtol=0, atol=0, check_dtype=True)
+
+
+@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
+def test_linear_backward_override_high_precision_forces_save_original_input(
+ recipe_name: str,
+) -> None:
+ reset_rng_states()
+ dtype = torch.bfloat16
+ input_shape = (32, 128)
+ _maybe_skip_recipe_dtype(recipe_name, dtype, "linear")
+ _maybe_skip_unsupported_recipe_module_combo(recipe_name, "linear")
+ _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, "linear")
+
+ mode_recipe = make_recipe(recipe_name, backward_override="high_precision")
+ skip_unsupported_backward_override("linear", mode_recipe, "high_precision")
+
+ module = te.Linear(
+ input_shape[-1],
+ 128,
+ bias=False,
+ params_dtype=dtype,
+ device="cuda",
+ save_original_input=False,
+ )
+ x = torch.randn(*input_shape, dtype=dtype, device="cuda")
+
+ _, _, saved_operands = _run_single_step_with_saved_operands(module, x, mode_recipe)
+
+ assert isinstance(saved_operands[0], torch.Tensor)
+
+
+@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
+def test_grouped_linear_backward_override_high_precision_forces_save_original_input(
+ recipe_name: str,
+) -> None:
+ reset_rng_states()
+ dtype = torch.bfloat16
+ in_features = 128
+ out_features = 128
+ m_splits = [64, 64]
+ num_gemms = len(m_splits)
+ num_tokens = sum(m_splits)
+ _maybe_skip_recipe_dtype(recipe_name, dtype, "grouped_linear")
+ _maybe_skip_unsupported_recipe_module_combo(recipe_name, "grouped_linear")
+ _maybe_skip_unsupported_grouped_splits(recipe_name, m_splits)
+
+ mode_recipe = make_recipe(recipe_name, backward_override="high_precision")
+ skip_unsupported_backward_override("grouped_linear", mode_recipe, "high_precision")
+
+ module = te.GroupedLinear(
+ num_gemms,
+ in_features,
+ out_features,
+ bias=False,
+ params_dtype=dtype,
+ device="cuda",
+ save_original_input=False,
+ )
+ x = torch.randn(num_tokens, in_features, dtype=dtype, device="cuda")
+
+ _, _, saved_operands = _run_grouped_linear_step_with_saved_operands(
+ module, x, m_splits, mode_recipe
+ )
+
+ saved_inputs = saved_operands[:num_gemms]
+ assert isinstance(saved_inputs[0], torch.Tensor)
+ assert saved_inputs[0].shape == x.shape
+ assert all(saved_input is None for saved_input in saved_inputs[1:])
+
+ saved_weights = saved_operands[2 * num_gemms : 3 * num_gemms]
+ for saved_weight in saved_weights:
+ assert isinstance(saved_weight, torch.Tensor)
+
+
@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
@pytest.mark.parametrize("module_type", ("linear", "layernorm_linear", "ops_linear"))
@pytest.mark.parametrize("input_shape,out_features", _shape_test_cases)
diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py
index ab12216df8..68d3ed9565 100644
--- a/tests/pytorch/test_fused_router.py
+++ b/tests/pytorch/test_fused_router.py
@@ -523,6 +523,76 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multipl
torch.testing.assert_close(probs.grad, probs_clone.grad, atol=atol, rtol=rtol)
+def test_fused_moe_aux_loss_cuda_graph_capture():
+ """CUDA-graph-safe path: total_num_tokens is a device tensor whose value
+ changes between replays. Forward and backward must both observe the new
+ value via the device-side coefficient computation."""
+ dtype = torch.float32
+ num_tokens = 4096
+ num_experts = 128
+ topk = 4
+ num_cols = num_experts
+ coeff = 0.01
+
+ offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4
+ probs = (
+ torch.arange(-num_cols // 2, num_cols // 2, device="cuda", dtype=dtype) * 1e-2
+ ).unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1)
+ probs = probs.contiguous().requires_grad_(True)
+ tokens_per_expert = torch.randint(1, 1000, (num_cols,), device="cuda", dtype=torch.int32)
+
+ total_num_tokens_dev = torch.tensor(num_tokens, dtype=torch.int64, device="cuda")
+
+ # Warmup on a side stream to satisfy CUDA Graph capture requirements.
+ s = torch.cuda.Stream()
+ s.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(s):
+ for _ in range(3):
+ warmup_out = fused_moe_aux_loss(
+ probs=probs,
+ tokens_per_expert=tokens_per_expert,
+ total_num_tokens=total_num_tokens_dev,
+ num_experts=num_experts,
+ topk=topk,
+ coeff=coeff,
+ )
+ torch.autograd.grad(warmup_out, probs)
+ del warmup_out
+ torch.cuda.current_stream().wait_stream(s)
+
+ g = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(g):
+ out = fused_moe_aux_loss(
+ probs=probs,
+ tokens_per_expert=tokens_per_expert,
+ total_num_tokens=total_num_tokens_dev,
+ num_experts=num_experts,
+ topk=topk,
+ coeff=coeff,
+ )
+ (grad_probs,) = torch.autograd.grad(out, probs)
+
+ atol, rtol = _get_tolerances(dtype, num_cols)
+ # Replay with several distinct token counts; the captured graph must pick
+ # up each new value through total_num_tokens_dev.
+ for new_total in (num_tokens, num_tokens // 2, num_tokens * 2 - 17):
+ total_num_tokens_dev.fill_(new_total)
+ g.replay()
+ torch.cuda.synchronize()
+ ref_probs = probs.detach().clone().requires_grad_(True)
+ ref = aux_loss_pytorch(
+ probs=ref_probs,
+ tokens_per_expert=tokens_per_expert,
+ total_num_tokens=new_total,
+ topk=topk,
+ num_experts=num_experts,
+ moe_aux_loss_coeff=coeff,
+ )
+ (ref_grad_probs,) = torch.autograd.grad(ref, ref_probs)
+ torch.testing.assert_close(out, ref, atol=atol, rtol=rtol)
+ torch.testing.assert_close(grad_probs, ref_grad_probs, atol=atol, rtol=rtol)
+
+
def _bytemap_to_bitmap_u8(bytemap: torch.Tensor) -> torch.Tensor:
"""Reference packer: bool[T, E] -> uint8[T, ceil(E/8)] LSB-first.
diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py
index caa84ec02a..01a7cf2415 100644
--- a/tests/pytorch/test_grouped_linear.py
+++ b/tests/pytorch/test_grouped_linear.py
@@ -1496,6 +1496,7 @@ def test_fp8_grouped_gemm(shape, accumulate):
_FUSED_GROUPED_GEMM_ENV = "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM"
_ALL_BOOLEAN = all_boolean
+_fp8_available, _reason_for_no_fp8 = fp8_available, reason_for_no_fp8
_mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8
_nvfp4_available, _reason_for_no_nvfp4 = nvfp4_available, reason_for_no_nvfp4
@@ -1577,6 +1578,10 @@ def _run_grouped_linear_path(
"fp8_recipe",
[
None,
+ pytest.param(
+ recipe.Float8CurrentScaling(),
+ marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8),
+ ),
pytest.param(
recipe.MXFP8BlockScaling(),
marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8),
@@ -1586,7 +1591,7 @@ def _run_grouped_linear_path(
marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4),
),
],
- ids=["bf16", "mxfp8", "nvfp4"],
+ ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"],
)
@pytest.mark.parametrize("bias", _ALL_BOOLEAN)
@pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN)
@@ -1600,8 +1605,13 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy(
pytest.skip(
"GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)."
)
- if use_fp8 and device_capability < (10, 0):
- pytest.skip("Quantized GroupedTensor grouped GEMM path requires Blackwell (SM100+).")
+ # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor
+ # current scaling also runs on the Hopper grouped GEMM path.
+ is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling()
+ if use_fp8 and not is_current_scaling and device_capability < (10, 0):
+ pytest.skip(
+ "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)."
+ )
cublaslt_version = tex.get_cublasLt_version()
if device_capability < (10, 0) and cublaslt_version < 130400:
pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
@@ -1786,6 +1796,10 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
"fp8_recipe",
[
None,
+ pytest.param(
+ recipe.Float8CurrentScaling(),
+ marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8),
+ ),
pytest.param(
recipe.MXFP8BlockScaling(),
marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8),
@@ -1795,7 +1809,7 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4),
),
],
- ids=["bf16", "mxfp8", "nvfp4"],
+ ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4"],
)
@pytest.mark.parametrize("bias", _ALL_BOOLEAN)
def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch):
@@ -1806,8 +1820,13 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch
pytest.skip(
"GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)."
)
- if use_fp8 and device_capability < (10, 0):
- pytest.skip("Quantized GroupedTensor grouped GEMM path requires Blackwell (SM100+).")
+ # MXFP8/NVFP4 grouped quantization kernels require Blackwell, but FP8 per-tensor
+ # current scaling also runs on the Hopper grouped GEMM path.
+ is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling()
+ if use_fp8 and not is_current_scaling and device_capability < (10, 0):
+ pytest.skip(
+ "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)."
+ )
cublaslt_version = tex.get_cublasLt_version()
if device_capability < (10, 0) and cublaslt_version < 130400:
pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index cb90ac6bd9..e24fff9049 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -288,13 +288,9 @@ def test_grouped_linear(
if single_grouped_bias and not bias:
pytest.skip("single_grouped_bias requires bias=True")
- if (
- single_grouped_weight
- and quantized_weight
- and quantization in ("fp8_delayed_scaling", "fp8_current_scaling")
- ):
+ if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"):
pytest.skip(
- "single_grouped_weight does not support FP8 delayed/current scaling "
+ "single_grouped_weight does not support FP8 delayed scaling "
"with quantized_model_init"
)
@@ -439,7 +435,10 @@ def test_grouped_linear(
@pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16))
@pytest.mark.parametrize(
"quantization",
- [None] + (["mxfp8"] if mxfp8_available else []),
+ [None]
+ + (["fp8_current_scaling"] if fp8_available else [])
+ + (["mxfp8"] if mxfp8_available else [])
+ + (["nvfp4_rht"] if nvfp4_available else []),
)
@pytest.mark.parametrize("quantized_weight", (False, True))
@pytest.mark.parametrize("bias", (False, True))
@@ -475,10 +474,38 @@ def test_grouped_linear_cuda_graph_safe(
"single_grouped_weight/single_grouped_bias requires"
" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1"
)
- if torch.cuda.get_device_capability() < (10, 0):
- pytest.skip("Grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)")
+ device_capability = torch.cuda.get_device_capability()
+ if device_capability < (9, 0):
+ pytest.skip(
+ "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)"
+ )
+ # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path,
+ # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+).
+ requires_blackwell = quantization is not None and quantization != "fp8_current_scaling"
+ if requires_blackwell and device_capability < (10, 0):
+ pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)")
+ # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+.
+ cublaslt_version = tex.get_cublasLt_version()
+ if device_capability < (10, 0) and cublaslt_version < 130400:
+ pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
+ if cublaslt_version < 130300:
+ pytest.skip("Grouped GEMM requires cuBLAS 13.3+.")
if quantization is None and quantized_weight:
pytest.skip("quantized_weight requires a quantization recipe")
+ if (
+ quantization is not None
+ and quantization.startswith("nvfp4")
+ and dtype != torch.bfloat16
+ ):
+ pytest.skip("NVFP4 grouped GEMM only supports BF16 output")
+ if single_grouped_weight and quantization is not None and quantization.startswith("nvfp4"):
+ # Currently, split_quantization is used which is not cuda graph safe.
+ # We should either support grouped weight quantization without rht or need to do
+ # inplace per tensor weight quantization to make this use-case cuda graphable if needed.
+ pytest.skip(
+ "NVFP4 grouped GEMM with single_grouped_weight is not supported yet; "
+ "only discrete weights (single_grouped_weight=False) are supported."
+ )
single_grouped_bias = bias and single_grouped_weight
diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py
index e1e15e6875..eeb6e7a394 100644
--- a/tests/pytorch/test_grouped_tensor.py
+++ b/tests/pytorch/test_grouped_tensor.py
@@ -33,6 +33,21 @@
mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True)
nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True)
+# The fused grouped FP8 block-scaling quantize/dequantize kernels are Hopper-only: they gate on
+# SM90-SM99 (NVTE_CHECK(sm >= 90 && sm < 100)). FP8 block scaling is still reported "available" on
+# Blackwell (SM100+) for the emulated/non-grouped paths, so ``fp8_block_scaling_available`` alone
+# does not exclude SM100 — add the Hopper arch bound for the grouped tests.
+_device_cc = torch.cuda.get_device_capability() if torch.cuda.is_available() else (0, 0)
+fp8_block_scaling_grouped_available = fp8_block_scaling_available and (9, 0) <= _device_cc < (10, 0)
+reason_for_no_fp8_block_scaling_grouped = (
+ reason_for_no_fp8_block_scaling
+ if not fp8_block_scaling_available
+ else (
+ "Fused grouped FP8 block-scaling quantize/dequantize is only supported on Hopper"
+ " (SM90-SM99)."
+ )
+)
+
_quantization_params = [
pytest.param(
"fp8_delayed_scaling",
@@ -115,6 +130,33 @@ def _rowwise_offset_bytes(numel: int, quantization: str) -> int:
return numel
+def _fp8bs_per_expert_scale_floats(
+ block_scaling_dim: int, columnwise: bool, m_t: int, k: int
+) -> int:
+ """Per-expert padded scale size (in floats) for grouped FP8 block-scaling.
+
+ Mirrors the per-expert sub-block layout that cuBLAS grouped GEMM consumes (and that the C++
+ test ``test_cast_float8blockwise_grouped.cu`` verifies against)::
+
+ 1D rowwise : blocks_X * roundup(M_t, 4)
+ 1D colwise : blocks_y_t * roundup(K, 4)
+ 2D rowwise : blocks_y_t * roundup(blocks_X, 4)
+ 2D colwise : blocks_X * roundup(blocks_y_t, 4)
+
+ The 2D columnwise roundup of each expert's block-rows to a multiple of 4 is the source of the
+ per-expert slack reserved in ``Float8BlockQuantizer.create_grouped_tensor``.
+ """
+
+ def align4(x: int) -> int:
+ return ((x + 3) // 4) * 4
+
+ blocks_x = (k + 127) // 128
+ blocks_y = (m_t + 127) // 128
+ if block_scaling_dim == 1:
+ return blocks_y * align4(k) if columnwise else blocks_x * align4(m_t)
+ return blocks_x * align4(blocks_y) if columnwise else blocks_y * align4(blocks_x)
+
+
class TestGroupedTensor:
@staticmethod
def setup_class(cls) -> None:
@@ -715,6 +757,54 @@ def _run_group_quantize(input_tensor):
if output_dbias:
assert torch.allclose(static_dbias, expected_dbias)
+ @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"])
+ @pytest.mark.skipif(
+ not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped
+ )
+ def test_group_quantize_fp8_blockwise_cudagraph_capturable(
+ self, block_scaling_dim: int
+ ) -> None:
+ """Ensure grouped FP8 block-scaling quantize is CUDA graph capturable (parity with MXFP8)."""
+ first_dims_host = [256, 128, 384]
+ num_tensors = len(first_dims_host)
+ hidden = 512
+ shape = [(r, hidden) for r in first_dims_host]
+ input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape]
+ grouped_input = torch.cat(input_tensors, dim=0)
+ first_dims = torch.tensor(first_dims_host, dtype=torch.int64, device="cuda")
+
+ quantizer = Float8BlockQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ rowwise=True,
+ columnwise=False,
+ force_pow_2_scales=False,
+ amax_epsilon=0.0,
+ block_scaling_dim=block_scaling_dim,
+ )
+
+ torch.cuda.synchronize()
+ static_input = grouped_input.clone()
+ static_first_dims = first_dims.clone()
+
+ def _run(inp):
+ return tex.group_quantize(inp, quantizer, num_tensors, static_first_dims)
+
+ _ = _run(static_input) # warmup allocator/kernels
+ torch.cuda.synchronize()
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ static_output = _run(static_input)
+
+ # Replay with fresh input copied into the captured buffer.
+ static_input.copy_(torch.randn_like(grouped_input))
+ graph.replay()
+ torch.cuda.synchronize()
+
+ expected = _run(static_input)
+ assert torch.equal(static_output.rowwise_data, expected.rowwise_data)
+ assert torch.equal(static_output.scale_inv, expected.scale_inv)
+
@pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"])
@pytest.mark.parametrize(
"shape_case",
@@ -914,6 +1004,117 @@ def _assert_fp8_cs_group_quantize_matches_reference(
expected = torch.cat(expected_columnwise)
assert torch.equal(grouped_output.columnwise_data[: expected.numel()], expected)
+ @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"])
+ @pytest.mark.parametrize("shape_case", ["uniform", "varying_first"])
+ @pytest.mark.parametrize("direction", ["rowwise", "columnwise", "both"])
+ @pytest.mark.parametrize("output_dbias", [False, True])
+ @pytest.mark.skipif(
+ not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped
+ )
+ def test_quantize_grouped_fp8_blockwise(
+ self, block_scaling_dim: int, shape_case: str, direction: str, output_dbias: bool
+ ) -> None:
+ """Test grouped FP8 block-scaling quantization against per-tensor quantization.
+
+ Covers rowwise, columnwise and both directions. Each expert's data sub-block (and, for
+ rowwise, its scale sub-block placed at the cumulative padded offset from
+ ``_fp8bs_per_expert_scale_floats``) is compared against an independent per-tensor reference
+ quantizer. The columnwise scale layout carries 2D padding plus the per-expert slack
+ reserved in ``Float8BlockQuantizer.create_grouped_tensor``; it is validated end to end by
+ ``test_group_dequantize_fp8_blockwise`` rather than by a fragile byte-compare here.
+
+ FP8 block-scaling supports only SAME_BOTH_DIMS (``uniform``) and VARYING_FIRST_DIM
+ (``varying_first``); ``varying_last``/``varying_both`` are rejected at the kernel level.
+ Per-tensor first dim must be a multiple of 128 (kernel tile size).
+ """
+ rowwise = direction in ("rowwise", "both")
+ columnwise = direction in ("columnwise", "both")
+
+ # dbias is the bias gradient (per-column input sum) emitted by the bgrad path, which
+ # requires rowwise output; the columnwise-only + dbias combination is not applicable.
+ if output_dbias and not rowwise:
+ pytest.skip("bgrad (dbias) requires rowwise output; columnwise-only does not apply.")
+
+ if shape_case == "uniform":
+ per_tensor_shapes = [(128, 512)] * 3
+ first_dims_host = None
+ else: # varying_first
+ per_tensor_shapes = [(128, 512), (256, 512), (384, 512)]
+ first_dims_host = [s[0] for s in per_tensor_shapes]
+
+ num_tensors = len(per_tensor_shapes)
+
+ input_tensors = [
+ torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in per_tensor_shapes
+ ]
+ flat_buffer = torch.cat([t.reshape(-1) for t in input_tensors])
+ common_last = per_tensor_shapes[0][1]
+ grouped_input = flat_buffer.view(-1, common_last)
+
+ first_dims = (
+ torch.tensor(first_dims_host, dtype=torch.int64, device="cuda")
+ if first_dims_host is not None
+ else None
+ )
+
+ quantizer = Float8BlockQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ rowwise=rowwise,
+ columnwise=columnwise,
+ force_pow_2_scales=False,
+ amax_epsilon=0.0,
+ block_scaling_dim=block_scaling_dim,
+ )
+
+ if output_dbias:
+ grouped_output, dbias = tex.bgrad_group_quantize(
+ grouped_input, quantizer, num_tensors, first_dims
+ )
+ else:
+ grouped_output = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims)
+
+ # Compare each expert's sub-block against an independent per-tensor reference. The
+ # reference enables both directions: the non-grouped 2D kernel requires rowwise output to
+ # be allocated even when only columnwise is consumed, and the columnwise data/scale it
+ # emits are independent of whether rowwise is also computed.
+ ref_quantizer = Float8BlockQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ rowwise=True,
+ columnwise=True,
+ force_pow_2_scales=False,
+ amax_epsilon=0.0,
+ block_scaling_dim=block_scaling_dim,
+ )
+ # Data sub-blocks are contiguous (rowwise (M_t, K) / columnwise transposed (K, M_t)) with
+ # no inter-expert padding. The rowwise scale sub-block is placed at the cumulative
+ # per-expert padded offset, so a wrong stride fails byte-equality here. The columnwise
+ # scale carries 2D ``roundup(blocks_y_t, 4)`` padding columns (and the per-expert slack),
+ # so its layout is validated end to end by ``test_group_dequantize_fp8_blockwise`` instead.
+ data_off = 0
+ rw_scale_off = 0
+ for tensor in input_tensors:
+ m_t, k = tensor.shape
+ numel = m_t * k
+ ref = ref_quantizer(tensor)
+ if rowwise:
+ ref_rw = ref._rowwise_data.reshape(-1)
+ assert torch.equal(grouped_output.rowwise_data[data_off : data_off + numel], ref_rw)
+ ref_rs = ref._rowwise_scale_inv.reshape(-1)
+ assert torch.equal(
+ grouped_output.scale_inv[rw_scale_off : rw_scale_off + ref_rs.numel()], ref_rs
+ )
+ rw_scale_off += _fp8bs_per_expert_scale_floats(block_scaling_dim, False, m_t, k)
+ if columnwise:
+ ref_cw = ref._columnwise_data.reshape(-1)
+ assert torch.equal(
+ grouped_output.columnwise_data[data_off : data_off + numel], ref_cw
+ )
+ data_off += numel
+
+ if output_dbias:
+ expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors])
+ assert torch.allclose(dbias, expected_dbias)
+
@pytest.mark.parametrize(
"shape",
[[(512, 1024), (512, 1024)], [(256, 512), (512, 512), (768, 512)]],
@@ -995,6 +1196,115 @@ def test_group_dequantize_cudagraph_capturable(self) -> None:
for exp, got in zip(expected_tensors, static_tensors):
assert torch.equal(got, exp)
+ @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"])
+ @pytest.mark.parametrize("direction", ["rowwise", "columnwise"])
+ @pytest.mark.skipif(
+ not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped
+ )
+ def test_group_dequantize_fp8_blockwise_cudagraph_capturable(
+ self, block_scaling_dim: int, direction: str
+ ) -> None:
+ """Ensure grouped FP8 block-scaling dequantize is CUDA graph capturable (parity with MXFP8)."""
+ rowwise = direction == "rowwise"
+ columnwise = direction == "columnwise"
+ num_tensors = 2
+ shape = [(512, 1024) for _ in range(num_tensors)]
+ input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape]
+ grouped_input = torch.cat(input_tensors, dim=0)
+
+ quantizer = Float8BlockQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ rowwise=rowwise,
+ columnwise=columnwise,
+ force_pow_2_scales=False,
+ amax_epsilon=0.0,
+ block_scaling_dim=block_scaling_dim,
+ )
+ first_dims = torch.tensor(
+ [shape[0][0] for _ in range(num_tensors)], dtype=torch.int64, device="cuda"
+ )
+
+ quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims)
+
+ # Warmup dequantize.
+ torch.cuda.synchronize()
+ _ = tex.group_dequantize(quantized, te.DType.kBFloat16)
+ torch.cuda.synchronize()
+
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ static_output = tex.group_dequantize(quantized, te.DType.kBFloat16)
+
+ # Replay with fresh quantized data copied into the captured input buffers.
+ fresh_input = torch.cat(
+ [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape], dim=0
+ )
+ fresh_quantized = tex.group_quantize(fresh_input, quantizer, num_tensors, first_dims)
+ if rowwise:
+ quantized.rowwise_data.copy_(fresh_quantized.rowwise_data)
+ quantized.scale_inv.copy_(fresh_quantized.scale_inv)
+ else:
+ quantized.columnwise_data.copy_(fresh_quantized.columnwise_data)
+ quantized.columnwise_scale_inv.copy_(fresh_quantized.columnwise_scale_inv)
+
+ graph.replay()
+ torch.cuda.synchronize()
+
+ expected = tex.group_dequantize(quantized, te.DType.kBFloat16)
+ expected_tensors = expected.split_into_quantized_tensors()
+ static_tensors = static_output.split_into_quantized_tensors()
+ for exp, got in zip(expected_tensors, static_tensors):
+ assert torch.equal(got, exp)
+
+ @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"])
+ @pytest.mark.parametrize("direction", ["rowwise", "columnwise"])
+ @pytest.mark.parametrize(
+ "shape",
+ [[(512, 1024), (512, 1024)], [(128, 512), (256, 512), (384, 512)]],
+ )
+ @pytest.mark.skipif(
+ not fp8_block_scaling_grouped_available, reason=reason_for_no_fp8_block_scaling_grouped
+ )
+ def test_group_dequantize_fp8_blockwise(
+ self, block_scaling_dim: int, direction: str, shape: List[Tuple[int, int]]
+ ) -> None:
+ """Test grouped FP8 block-scaling dequantize round-trip for rowwise and columnwise.
+
+ The columnwise + ``varying_first`` + 2D case exercises the per-expert columnwise
+ scale-buffer slack end to end: a wrong slack/stride would place scales at the wrong
+ offsets and corrupt the dequantized values. ``group_dequantize`` consumes exactly one of
+ rowwise / columnwise data, so ``both`` is not a valid round-trip and is not tested here.
+ """
+ rowwise = direction == "rowwise"
+ columnwise = direction == "columnwise"
+ num_tensors = len(shape)
+
+ input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape]
+ grouped_input = torch.cat(input_tensors, dim=0)
+
+ quantizer = Float8BlockQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ rowwise=rowwise,
+ columnwise=columnwise,
+ force_pow_2_scales=False,
+ amax_epsilon=0.0,
+ block_scaling_dim=block_scaling_dim,
+ )
+ first_dims = torch.tensor([s[0] for s in shape], dtype=torch.int64, device="cuda")
+
+ quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims)
+ dequantized = tex.group_dequantize(quantized, te.DType.kBFloat16)
+
+ assert dequantized.num_tensors == num_tensors
+ assert dequantized.logical_shape == quantized.logical_shape
+ assert torch.equal(dequantized.first_dims, quantized.first_dims)
+ assert torch.equal(dequantized.tensor_offsets, quantized.tensor_offsets)
+
+ dequantized_tensors = dequantized.split_into_quantized_tensors()
+ assert len(dequantized_tensors) == num_tensors
+ for orig, deq in zip(input_tensors, dequantized_tensors):
+ torch.testing.assert_close(deq, orig, atol=0.125, rtol=0.1)
+
def test_clear(self) -> None:
"""Test clear method"""
num_tensors = 3
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 1286492a6e..8ebce563ce 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -32,6 +32,9 @@
is_mxfp8_available,
is_fp8_block_scaling_available,
is_nvfp4_available,
+ Float8BlockQuantizer,
+ MXFP8Quantizer,
+ NVFP4Quantizer,
)
from utils import recipe_id
@@ -384,3 +387,170 @@ def fn(inp):
out = compiled(inp)
out.sum().backward()
+
+
+# ---------------------------------------------------------------------------
+# Value-opaque quantizers
+# ---------------------------------------------------------------------------
+
+
+def _mxfp8(dtype=tex.DType.kFloat8E4M3):
+ return MXFP8Quantizer(fp8_dtype=dtype)
+
+
+def _blockwise(force_pow_2_scales=True):
+ return Float8BlockQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ rowwise=True,
+ columnwise=True,
+ force_pow_2_scales=force_pow_2_scales,
+ )
+
+
+def _current_scaling(amax_epsilon=0.0):
+ return Float8CurrentScalingQuantizer(
+ fp8_dtype=tex.DType.kFloat8E4M3,
+ device=torch.device("cpu"),
+ amax_epsilon=amax_epsilon,
+ )
+
+
+def _nvfp4(with_rht=True):
+ # Default with_rht=True so the quantize round-trip below exercises the
+ # derived ``rht_matrix`` tensor (the field most likely to be dropped on
+ # value-key reconstruction). Post-RHT amax is required by the kernel
+ # whenever RHT is on (pre-RHT amax is unsupported).
+ return NVFP4Quantizer(
+ fp4_dtype=tex.DType.kFloat4E2M1,
+ rowwise=True,
+ columnwise=True,
+ with_rht=with_rht,
+ with_post_rht_amax=with_rht,
+ )
+
+
+def _hw_available(quantizer):
+ """Whether this HW can actually run the quantize kernel for *quantizer*."""
+ if isinstance(quantizer, MXFP8Quantizer):
+ return mxfp8_available
+ if isinstance(quantizer, NVFP4Quantizer):
+ return nvfp4_available
+ if isinstance(quantizer, Float8BlockQuantizer):
+ return fp8_block_scaling_available
+ return fp8_available # Float8CurrentScalingQuantizer
+
+
+# (factory, kwargs producing a different-but-valid config)
+_VALUE_QUANTIZERS = [
+ pytest.param(_mxfp8, id="mxfp8"),
+ pytest.param(_blockwise, id="float8_blockwise"),
+ pytest.param(_current_scaling, id="float8_current_scaling"),
+ pytest.param(
+ _nvfp4,
+ id="nvfp4",
+ marks=pytest.mark.skipif(
+ not torch.cuda.is_available(),
+ reason="NVFP4Quantizer requires CUDA to construct",
+ ),
+ ),
+]
+
+
+@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS)
+def test_quantizer_value_object(factory):
+ """Value semantics + ``__fx_repr__`` round-trip via the production FX path."""
+ a = factory()
+
+ # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object.
+ repr_str, globals_ = a.__fx_repr__()
+ rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used
+ assert rebuilt == a and rebuilt is not a
+ assert hash(rebuilt) == hash(a)
+
+ # The rebuilt quantizer must also *behave* identically, not just compare
+ # equal: equality only looks at the value key, so a field the kernel needs
+ # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would
+ # slip through the checks above and only blow up at quantize time. Run the
+ # real quantize kernel on both and require bit-exact results.
+ if torch.cuda.is_available() and _hw_available(a):
+ x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+ torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0)
+
+
+def test_value_quantizer_rejects_process_group():
+ """A value quantizer holding a live ProcessGroup must refuse to be turned
+ into a value key / FX constant (raise), not silently drop the group."""
+ import torch.distributed as dist # pylint: disable=import-outside-toplevel
+
+ created = not dist.is_initialized()
+ if created:
+ dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1)
+ try:
+ q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)
+ q.amax_reduction_group = dist.group.WORLD
+ # Every value-materialization path must reject it (hash, eq, __fx_repr__).
+ with pytest.raises(TypeError):
+ hash(q)
+ with pytest.raises(TypeError):
+ q.__fx_repr__()
+ finally:
+ if created:
+ dist.destroy_process_group()
+
+
+if _opaque_available:
+ # A minimal custom op taking a tensor and a value-opaque quantizer that
+ # quantizes + dequantizes inside it, one per production quantizer class.
+ # ``test_quantizer_value_object_fullgraph`` drives this under
+ # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the
+ # graph -- proving the opaque-type registration took effect (a graph break
+ # would make ``fullgraph=True`` raise).
+ _qdq_lib = torch.library.Library("test_te_qdq", "DEF")
+ _QDQ_OPS = {}
+ for _qcls in (
+ MXFP8Quantizer,
+ Float8BlockQuantizer,
+ Float8CurrentScalingQuantizer,
+ NVFP4Quantizer,
+ ):
+ _op = f"qdq_{_qcls.__name__}"
+ _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor")
+
+ @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib)
+ def _qdq_impl(x, q):
+ return q(x).dequantize()
+
+ @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib)
+ def _qdq_fake(x, q):
+ return torch.empty_like(x)
+
+ _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op)
+
+
+@pytest.mark.skipif(
+ not _opaque_available,
+ reason="torch.compile opaque-object support requires PyTorch >= 2.11",
+)
+@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS)
+def test_quantizer_value_object_fullgraph(factory):
+ """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph.
+
+ A custom op quantizes+dequantizes with the (opaque value) quantizer; the
+ compiled result must match eager. ``fullgraph=True`` raises on any graph
+ break, so this proves the opaque-type registration actually took effect --
+ unlike merely passing the quantizer through.
+ """
+ q = factory()
+ if not (torch.cuda.is_available() and _hw_available(q)):
+ pytest.skip("format not supported on this HW")
+
+ op = _QDQ_OPS[type(q)]
+ x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+
+ def fn(inp):
+ return op(inp, q)
+
+ ref = fn(x)
+ torch._dynamo.reset()
+ out = torch.compile(fn, fullgraph=True)(x)
+ torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh
index 63c1b046ff..bf4a021811 100644
--- a/transformer_engine/common/cast/dispatch/dequantize.cuh
+++ b/transformer_engine/common/cast/dispatch/dequantize.cuh
@@ -15,6 +15,7 @@
#include "../../common.h"
#include "../fp8/dequantize_fp8.cuh"
+#include "../fp8_blockwise/group_dequantize_fp8_blockwise.cuh"
#include "../mxfp8/dequantize_mxfp8.cuh"
#include "../mxfp8/group_dequantize_mxfp8.cuh"
#include "../nvfp4/dequantize_nvfp4.cuh"
@@ -69,6 +70,11 @@ inline void group_dequantize_helper(const GroupedTensor &input, GroupedTensor *o
}
break;
}
+ case NVTE_BLOCK_SCALING_1D:
+ case NVTE_BLOCK_SCALING_2D: {
+ fp8_blockwise::group_dequantize(&input, output, stream);
+ break;
+ }
default:
NVTE_ERROR("Grouped dequantize not implemented for scaling mode: " +
to_string(input.scaling_mode) + ".");
diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh
index 97dd27aec6..031122d966 100644
--- a/transformer_engine/common/cast/dispatch/quantize.cuh
+++ b/transformer_engine/common/cast/dispatch/quantize.cuh
@@ -19,6 +19,7 @@
#include "../core/common.cuh"
#include "../fp8/group_quantize_fp8.cuh"
#include "../fp8/quantize_fp8.cuh"
+#include "../fp8_blockwise/group_quantize_fp8_blockwise.cuh"
#include "../mxfp8/group_quantize_mxfp8.cuh"
#include "../mxfp8/quantize_mxfp8.cuh"
#include "../nvfp4/group_quantize_transpose_nvfp4.cuh"
@@ -472,6 +473,26 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor
workspace_tensor, &quant_config_cpp, stream);
break;
}
+ case NVTE_BLOCK_SCALING_1D: {
+ NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_1D.");
+ NVTE_CHECK(!quant_config_cpp.force_pow_2_scales,
+ "Fused grouped FP8 block-scaling quantize does not support "
+ "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused "
+ "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0).");
+ fp8_blockwise::group_quantize_blockwise_1d(input_tensor, output_tensor, noop_tensor,
+ quant_config_cpp.amax_epsilon, stream);
+ break;
+ }
+ case NVTE_BLOCK_SCALING_2D: {
+ NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_BLOCK_SCALING_2D.");
+ NVTE_CHECK(!quant_config_cpp.force_pow_2_scales,
+ "Fused grouped FP8 block-scaling quantize does not support "
+ "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused "
+ "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0).");
+ fp8_blockwise::group_quantize_blockwise_2d(input_tensor, output_tensor, noop_tensor,
+ quant_config_cpp.amax_epsilon, stream);
+ break;
+ }
default:
NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + ".");
}
@@ -513,6 +534,28 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe
&quant_config_cpp, stream);
break;
}
+ case NVTE_BLOCK_SCALING_1D:
+ case NVTE_BLOCK_SCALING_2D: {
+ NVTE_CHECK(!IS_DACT, "IS_DACT is not implemented for grouped FP8 block scaling.");
+ NVTE_CHECK(!quant_config_cpp.force_pow_2_scales,
+ "Fused grouped FP8 block-scaling quantize does not support "
+ "force_pow_2_scales=True. Set force_pow_2_scales=False, or use the unfused "
+ "split-quantize path (NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=0).");
+ // dbias is computed in-kernel and reduced per-expert inside group_quantize_blockwise_{1d,2d}
+ // (mirrors MXFP8); those also handle the two-call workspace sizing protocol.
+ GroupedTensor *dbias_arg = IS_DBIAS ? dbias_tensor : nullptr;
+ Tensor *workspace_arg = IS_DBIAS ? workspace_tensor : nullptr;
+ if (scaling_mode == NVTE_BLOCK_SCALING_1D) {
+ fp8_blockwise::group_quantize_blockwise_1d(grad_tensor, output_tensor, noop_tensor,
+ quant_config_cpp.amax_epsilon, stream, dbias_arg,
+ workspace_arg);
+ } else {
+ fp8_blockwise::group_quantize_blockwise_2d(grad_tensor, output_tensor, noop_tensor,
+ quant_config_cpp.amax_epsilon, stream, dbias_arg,
+ workspace_arg);
+ }
+ break;
+ }
default:
NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + ".");
}
diff --git a/transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh
new file mode 100644
index 0000000000..556dcad428
--- /dev/null
+++ b/transformer_engine/common/cast/fp8_blockwise/group_dequantize_fp8_blockwise.cuh
@@ -0,0 +1,519 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+/*! \file group_dequantize_fp8_blockwise.cuh
+ * \brief CUDA kernels to dequantize grouped tensors from FP8 with 1D/2D
+ * block scaling (rowwise or columnwise) back to BF16 / FP16 / FP32. Mirrors
+ * the per-expert layouts written by ``group_quantize_fp8_blockwise``.
+ */
+
+#ifndef TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_FP8_BLOCKWISE_CUH_
+#define TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_FP8_BLOCKWISE_CUH_
+
+#include
+#include
+#include
+
+#include "../../common.h"
+#include "../../utils.cuh"
+#include "../core/common.cuh"
+#include "group_quantize_fp8_blockwise.cuh"
+
+namespace transformer_engine {
+namespace dispatch {
+namespace fp8_blockwise {
+
+namespace group_dequantize_kernel {
+
+// Resolve which expert a tile (blocks_X x total_row_blocks grid) belongs to and its row range.
+// tensor_M (the expert's first-dim length) addresses the per-expert (K, M_t) transposed block
+// in the columnwise modes.
+struct TileExpertInfo {
+ size_t tensor_id;
+ size_t tensor_block_y_base;
+ size_t tensor_row_blocks;
+ size_t tensor_row_base;
+ size_t tensor_M;
+ bool in_bounds;
+};
+
+template
+__device__ __forceinline__ TileExpertInfo resolve_tile_expert(
+ size_t tile_y_global, size_t num_tensors, size_t common_first_dim_blocks, size_t K,
+ size_t total_row_blocks, const int64_t* __restrict__ tensor_offsets_ptr) {
+ TileExpertInfo info{};
+ info.in_bounds = false;
+ if (tile_y_global >= total_row_blocks) return info;
+ const size_t tile_row_stride = static_cast(kTileDim) * K;
+ info.tensor_id = find_tensor_id_by_block_y(
+ tile_y_global, num_tensors, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr);
+ info.tensor_block_y_base =
+ kSameBothDims
+ ? (info.tensor_id * common_first_dim_blocks)
+ : tensor_block_y_base_from_offsets(info.tensor_id, tensor_offsets_ptr, tile_row_stride);
+ info.tensor_row_blocks = kSameBothDims
+ ? common_first_dim_blocks
+ : (tensor_block_y_base_from_offsets(
+ info.tensor_id + 1, tensor_offsets_ptr, tile_row_stride) -
+ info.tensor_block_y_base);
+ if (tile_y_global >= info.tensor_block_y_base + info.tensor_row_blocks) return info;
+ info.tensor_row_base = info.tensor_block_y_base * kTileDim;
+ info.tensor_M = info.tensor_row_blocks * kTileDim;
+ info.in_bounds = true;
+ return info;
+}
+
+// ===== 1D rowwise =====
+// Per-expert scale layout: (blocks_X, roundup(M_t, 4)) floats.
+// scale[expert_off + tile_x * roundup(M_t, 4) + r_local]
+template
+__global__ void __launch_bounds__(kThreadsPerBlock)
+ group_dequantize_blockwise_1d_rw_kernel(const IType* __restrict__ input_base,
+ OType* __restrict__ output_base,
+ const CType* __restrict__ scale_inv_base,
+ const int64_t* __restrict__ tensor_offsets_ptr,
+ const size_t num_tensors,
+ const size_t common_first_dim_blocks, const size_t K,
+ const size_t total_row_blocks, const size_t R_total) {
+#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000
+ const size_t tile_x = blockIdx.x;
+ const size_t tile_y_global = blockIdx.y;
+ const auto info = resolve_tile_expert(
+ tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr);
+ if (!info.in_bounds) return;
+
+ const size_t blocks_X = DIVUP(K, static_cast(kTileDim));
+ const size_t tile_row_stride = static_cast(kTileDim) * K;
+ const size_t expert_offset = expert_scale_offset_1d_rowwise(
+ info.tensor_id, blocks_X, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr);
+ const size_t per_expert_stride = DIVUP_TO_MULTIPLE(info.tensor_M, kScaleColAlign);
+ const CType* const tile_scale_inv_base =
+ scale_inv_base + expert_offset + tile_x * per_expert_stride;
+
+ const size_t global_row_base = tile_y_global * kTileDim;
+ const size_t global_col_base = tile_x * kTileDim;
+
+ constexpr int kThreadsPerRow = 8;
+ constexpr int kVec = 16;
+ constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow; // 32
+ constexpr int kIters = kTileDim / kRowsPerIter; // 4
+
+ const int tid = threadIdx.x;
+ const int thr_col = tid % kThreadsPerRow;
+ const int thr_row = tid / kThreadsPerRow;
+ const size_t c = global_col_base + static_cast(thr_col) * kVec;
+
+#pragma unroll
+ for (int it = 0; it < kIters; ++it) {
+ const int row_local = thr_row + it * kRowsPerIter;
+ const size_t r_global = global_row_base + row_local;
+ if (r_global >= R_total) continue;
+
+ const size_t r_local = r_global - info.tensor_row_base;
+ const CType s_inv = tile_scale_inv_base[r_local];
+
+ Vec in_vec;
+ if (c + kVec <= K) {
+ in_vec.load_from(input_base + r_global * K + c);
+ } else if (c < K) {
+ in_vec.load_from_elts(input_base + r_global * K + c, 0, K - c);
+ } else {
+ continue;
+ }
+
+ Vec out_vec;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv);
+ }
+
+ if (c + kVec <= K) {
+ out_vec.store_to(output_base + r_global * K + c);
+ } else if (c < K) {
+ out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c);
+ }
+ }
+#endif
+}
+
+// ===== 1D columnwise =====
+// Data layout per-expert: (K, M_t) transposed -- element (r, c) is at
+// input[tensor_row_base * K + c * tensor_M + r_local].
+// Scale layout GLOBAL: (total_row_blocks, roundup(K, 4)) floats.
+// scale_inv[tile_y_global * scale_t_stride_aligned_K + c]
+template
+__global__ void __launch_bounds__(kThreadsPerBlock)
+ group_dequantize_blockwise_1d_cw_kernel(const IType* __restrict__ input_base,
+ OType* __restrict__ output_base,
+ const CType* __restrict__ scale_inv_base,
+ const size_t scale_t_stride_aligned_K,
+ const int64_t* __restrict__ tensor_offsets_ptr,
+ const size_t num_tensors,
+ const size_t common_first_dim_blocks, const size_t K,
+ const size_t total_row_blocks, const size_t R_total) {
+#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000
+ const size_t tile_x = blockIdx.x;
+ const size_t tile_y_global = blockIdx.y;
+ const auto info = resolve_tile_expert(
+ tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr);
+ if (!info.in_bounds) return;
+
+ const size_t expert_data_off = info.tensor_row_base * K;
+ const CType* const tile_scale_base = scale_inv_base + tile_y_global * scale_t_stride_aligned_K;
+
+ const size_t global_row_base = tile_y_global * kTileDim;
+ const size_t global_col_base = tile_x * kTileDim;
+
+ constexpr int kThreadsPerRow = 8;
+ constexpr int kVec = 16;
+ constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow;
+ constexpr int kIters = kTileDim / kRowsPerIter;
+
+ const int tid = threadIdx.x;
+ const int thr_col = tid % kThreadsPerRow;
+ const int thr_row = tid / kThreadsPerRow;
+ const size_t c = global_col_base + static_cast(thr_col) * kVec;
+
+ // 1D columnwise has one scale per column. Pre-load this thread's 16 columns.
+ CType s_inv[kVec];
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ s_inv[e] = (c + e < K) ? tile_scale_base[c + e] : static_cast(0.f);
+ }
+
+#pragma unroll
+ for (int it = 0; it < kIters; ++it) {
+ const int row_local = thr_row + it * kRowsPerIter;
+ const size_t r_global = global_row_base + row_local;
+ if (r_global >= R_total) continue;
+
+ const size_t r_local = r_global - info.tensor_row_base;
+
+ // Per-expert (K, M_t) transposed input: strided by M_t per column (no vector load).
+ // K % 128 == 0 so c+e < K always holds; the explicit else just keeps correctness
+ // independent of that invariant.
+ Vec in_vec;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ in_vec.data.elt[e] = (c + e < K)
+ ? input_base[expert_data_off + (c + e) * info.tensor_M + r_local]
+ : static_cast(0);
+ }
+
+ Vec out_vec;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv[e]);
+ }
+
+ if (c + kVec <= K) {
+ out_vec.store_to(output_base + r_global * K + c);
+ } else if (c < K) {
+ out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c);
+ }
+ }
+#endif
+}
+
+// ===== 2D rowwise =====
+// Data layout: (M, K) flat. One scale per 128x128 tile.
+// Scale layout GLOBAL: (total_row_blocks, roundup(blocks_X, 4)) floats.
+// scale[tile_y_global * scale_stride_y + tile_x]
+template
+__global__ void __launch_bounds__(kThreadsPerBlock)
+ group_dequantize_blockwise_2d_rw_kernel(const IType* __restrict__ input_base,
+ OType* __restrict__ output_base,
+ const CType* __restrict__ scale_inv_base,
+ const size_t scale_stride_y,
+ const int64_t* __restrict__ tensor_offsets_ptr,
+ const size_t num_tensors,
+ const size_t common_first_dim_blocks, const size_t K,
+ const size_t total_row_blocks, const size_t R_total) {
+#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000
+ const size_t tile_x = blockIdx.x;
+ const size_t tile_y_global = blockIdx.y;
+ const auto info = resolve_tile_expert(
+ tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr);
+ if (!info.in_bounds) return;
+
+ // 2D: one scale per tile.
+ const CType s_inv = scale_inv_base[tile_y_global * scale_stride_y + tile_x];
+
+ const size_t global_row_base = tile_y_global * kTileDim;
+ const size_t global_col_base = tile_x * kTileDim;
+
+ constexpr int kThreadsPerRow = 8;
+ constexpr int kVec = 16;
+ constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow;
+ constexpr int kIters = kTileDim / kRowsPerIter;
+
+ const int tid = threadIdx.x;
+ const int thr_col = tid % kThreadsPerRow;
+ const int thr_row = tid / kThreadsPerRow;
+ const size_t c = global_col_base + static_cast(thr_col) * kVec;
+
+#pragma unroll
+ for (int it = 0; it < kIters; ++it) {
+ const int row_local = thr_row + it * kRowsPerIter;
+ const size_t r_global = global_row_base + row_local;
+ if (r_global >= R_total) continue;
+
+ Vec in_vec;
+ if (c + kVec <= K) {
+ in_vec.load_from(input_base + r_global * K + c);
+ } else if (c < K) {
+ in_vec.load_from_elts(input_base + r_global * K + c, 0, K - c);
+ } else {
+ continue;
+ }
+
+ Vec out_vec;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv);
+ }
+
+ if (c + kVec <= K) {
+ out_vec.store_to(output_base + r_global * K + c);
+ } else if (c < K) {
+ out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c);
+ }
+ }
+#endif
+}
+
+// ===== 2D columnwise =====
+// Data layout per-expert: (K, M_t) transposed.
+// Scale layout per-expert: (blocks_X, roundup(blocks_y_t, 4)) floats.
+// scale[expert_off + tile_x * roundup(blocks_y_t, 4) + local_tile_y]
+template
+__global__ void __launch_bounds__(kThreadsPerBlock)
+ group_dequantize_blockwise_2d_cw_kernel(const IType* __restrict__ input_base,
+ OType* __restrict__ output_base,
+ const CType* __restrict__ scale_inv_base,
+ const int64_t* __restrict__ tensor_offsets_ptr,
+ const size_t num_tensors,
+ const size_t common_first_dim_blocks, const size_t K,
+ const size_t total_row_blocks, const size_t R_total) {
+#if __CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000
+ const size_t tile_x = blockIdx.x;
+ const size_t tile_y_global = blockIdx.y;
+ const auto info = resolve_tile_expert(
+ tile_y_global, num_tensors, common_first_dim_blocks, K, total_row_blocks, tensor_offsets_ptr);
+ if (!info.in_bounds) return;
+
+ // `info.in_bounds` is derived from blockIdx.y and is uniform across the CTA,
+ // so the early return above never strands a sibling thread inside
+ // compute_2d_cw_expert_offset's __syncthreads().
+ __shared__ size_t warp_offset_partials[kNumWarps];
+
+ const int tid = threadIdx.x;
+ const int warp_id = tid / kThreadsPerWarp;
+ const int lane = tid % kThreadsPerWarp;
+
+ const size_t blocks_X = DIVUP(K, static_cast(kTileDim));
+ const size_t tile_row_stride = static_cast(kTileDim) * K;
+ const size_t expert_offset = compute_2d_cw_expert_offset(
+ info.tensor_id, blocks_X, common_first_dim_blocks, tile_row_stride, tensor_offsets_ptr,
+ warp_offset_partials, tid, warp_id, lane);
+ const size_t per_expert_stride_t = DIVUP_TO_MULTIPLE(info.tensor_row_blocks, kScaleColAlign);
+ const size_t local_tile_y = tile_y_global - info.tensor_block_y_base;
+ const CType s_inv = scale_inv_base[expert_offset + tile_x * per_expert_stride_t + local_tile_y];
+
+ const size_t expert_data_off = info.tensor_row_base * K;
+ const size_t global_row_base = tile_y_global * kTileDim;
+ const size_t global_col_base = tile_x * kTileDim;
+
+ constexpr int kThreadsPerRow = 8;
+ constexpr int kVec = 16;
+ constexpr int kRowsPerIter = kThreadsPerBlock / kThreadsPerRow;
+ constexpr int kIters = kTileDim / kRowsPerIter;
+
+ const int thr_col = tid % kThreadsPerRow;
+ const int thr_row = tid / kThreadsPerRow;
+ const size_t c = global_col_base + static_cast(thr_col) * kVec;
+
+#pragma unroll
+ for (int it = 0; it < kIters; ++it) {
+ const int row_local = thr_row + it * kRowsPerIter;
+ const size_t r_global = global_row_base + row_local;
+ if (r_global >= R_total) continue;
+
+ const size_t r_local = r_global - info.tensor_row_base;
+
+ // Explicit else zeroes out-of-range lanes (c+e < K always holds since K % 128 == 0;
+ // this keeps correctness independent of that invariant).
+ Vec in_vec;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ in_vec.data.elt[e] = (c + e < K)
+ ? input_base[expert_data_off + (c + e) * info.tensor_M + r_local]
+ : static_cast(0);
+ }
+
+ Vec out_vec;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ out_vec.data.elt[e] = static_cast(static_cast(in_vec.data.elt[e]) * s_inv);
+ }
+
+ if (c + kVec <= K) {
+ out_vec.store_to(output_base + r_global * K + c);
+ } else if (c < K) {
+ out_vec.store_to_elts(output_base + r_global * K + c, 0, K - c);
+ }
+ }
+#endif
+}
+
+} // namespace group_dequantize_kernel
+
+// Host-side dispatcher. Supports all four combinations of {1D, 2D} block
+// scaling x {rowwise, columnwise} data, matching the layouts written by
+// ``group_quantize_fp8_blockwise``. The input GroupedTensor must have exactly
+// one of rowwise / columnwise data populated (the dequantize API rejects
+// both).
+inline void group_dequantize(const GroupedTensor* input, GroupedTensor* output,
+ cudaStream_t stream) {
+ using namespace group_dequantize_kernel;
+
+ const int sm = transformer_engine::cuda::sm_arch();
+ NVTE_CHECK(sm >= 90 && sm < 100,
+ "Grouped FP8 block-scaling dequantize is only supported on Hopper (SM90-SM99); "
+ "got SM",
+ sm, ".");
+ NVTE_CHECK(
+ input->scaling_mode == NVTE_BLOCK_SCALING_1D || input->scaling_mode == NVTE_BLOCK_SCALING_2D,
+ "Grouped FP8 block-scaling dequantize requires 1D or 2D block scaling "
+ "(got scaling_mode=",
+ to_string(input->scaling_mode), ").");
+ NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input must have FP8 type.");
+ NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision.");
+ NVTE_CHECK(!is_fp4_dtype(output->dtype()), "Output must not be FP4.");
+ NVTE_CHECK(input->num_tensors == output->num_tensors,
+ "Number of input and output tensors must match.");
+
+ const bool use_rowwise = input->has_data();
+ const bool use_colwise = input->has_columnwise_data();
+ NVTE_CHECK(use_rowwise || use_colwise, "Input must have rowwise or columnwise data populated.");
+ NVTE_CHECK(!(use_rowwise && use_colwise),
+ "Grouped FP8 block-scaling dequantize accepts exactly one direction at a "
+ "time (not both rowwise and columnwise simultaneously).");
+ NVTE_CHECK(!input->with_gemm_swizzled_scales,
+ "Grouped FP8 block-scaling dequantize requires compact (un-swizzled) scales.");
+
+ const size_t first_logical_dim = input->logical_shape.data[0];
+ const size_t last_logical_dim = input->logical_shape.data[1];
+ if (first_logical_dim == 0 || last_logical_dim == 0) return;
+
+ const bool same_both_dims = input->all_same_shape();
+ const bool varying_first_dim = (!input->all_same_first_dim()) && input->all_same_last_dim();
+ NVTE_CHECK(same_both_dims || varying_first_dim,
+ "Grouped FP8 block-scaling dequantize supports only SAME_BOTH_DIMS and "
+ "VARYING_FIRST_DIM shape representations.");
+
+ const size_t num_tensors = input->num_tensors;
+ const size_t K = last_logical_dim;
+ NVTE_CHECK(K % kTileDim == 0,
+ "Last dim must be a multiple of 128 for FP8 block-scaling dequantize (got ", K, ").");
+
+ size_t common_first_dim_blocks = 0;
+ if (same_both_dims) {
+ const size_t common_first_dim = input->get_common_first_dim();
+ NVTE_CHECK(common_first_dim % kTileDim == 0,
+ "SAME_BOTH_DIMS first dim must be multiple of 128 (got ", common_first_dim, ").");
+ common_first_dim_blocks = common_first_dim / kTileDim;
+ }
+ const size_t total_row_blocks = DIVUP(first_logical_dim, static_cast(kTileDim));
+ const size_t blocks_X = K / kTileDim;
+
+ const int64_t* tensor_offsets_ptr =
+ same_both_dims ? nullptr : reinterpret_cast(input->tensor_offsets.dptr);
+ if (!same_both_dims) {
+ NVTE_CHECK(tensor_offsets_ptr != nullptr,
+ "VARYING_FIRST_DIM requires tensor_offsets to be set on the input.");
+ }
+
+ const dim3 grid(blocks_X, total_row_blocks);
+ const dim3 block(kThreadsPerBlock);
+ const bool is_1d = (input->scaling_mode == NVTE_BLOCK_SCALING_1D);
+
+ // Pick the populated direction's data + scale buffer. The global-layout strides (1D cw, 2D rw)
+ // are derived from K / blocks_X exactly as the quantize launcher does.
+ const SimpleTensor& input_data = use_rowwise ? input->data : input->columnwise_data;
+ const SimpleTensor& input_scale_inv =
+ use_rowwise ? input->scale_inv : input->columnwise_scale_inv;
+ const size_t scale_t_stride_aligned_K = DIVUP_TO_MULTIPLE(K, kScaleColAlign);
+ const size_t scale_stride_y = DIVUP_TO_MULTIPLE(blocks_X, kScaleColAlign);
+
+ TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(
+ input->dtype(), IType,
+ TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
+ output->dtype(), OType, using CType = float;
+ const IType* const input_dptr = reinterpret_cast(input_data.dptr);
+ OType* const output_dptr = reinterpret_cast(output->data.dptr);
+ const CType* const scale_inv_dptr = reinterpret_cast(input_scale_inv.dptr);
+
+ if (is_1d && use_rowwise) {
+ if (same_both_dims) {
+ group_dequantize_blockwise_1d_rw_kernel
+ <<>>(
+ input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors,
+ common_first_dim_blocks, K, total_row_blocks, first_logical_dim);
+ } else {
+ group_dequantize_blockwise_1d_rw_kernel
+ <<>>(
+ input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors,
+ common_first_dim_blocks, K, total_row_blocks, first_logical_dim);
+ }
+ } else if (is_1d && use_colwise) {
+ if (same_both_dims) {
+ group_dequantize_blockwise_1d_cw_kernel
+ <<>>(input_dptr, output_dptr, scale_inv_dptr,
+ scale_t_stride_aligned_K, tensor_offsets_ptr,
+ num_tensors, common_first_dim_blocks, K,
+ total_row_blocks, first_logical_dim);
+ } else {
+ group_dequantize_blockwise_1d_cw_kernel
+ <<>>(input_dptr, output_dptr, scale_inv_dptr,
+ scale_t_stride_aligned_K, tensor_offsets_ptr,
+ num_tensors, common_first_dim_blocks, K,
+ total_row_blocks, first_logical_dim);
+ }
+ } else if (!is_1d && use_rowwise) {
+ if (same_both_dims) {
+ group_dequantize_blockwise_2d_rw_kernel
+ <<>>(
+ input_dptr, output_dptr, scale_inv_dptr, scale_stride_y, tensor_offsets_ptr,
+ num_tensors, common_first_dim_blocks, K, total_row_blocks, first_logical_dim);
+ } else {
+ group_dequantize_blockwise_2d_rw_kernel
+ <<>>(
+ input_dptr, output_dptr, scale_inv_dptr, scale_stride_y, tensor_offsets_ptr,
+ num_tensors, common_first_dim_blocks, K, total_row_blocks, first_logical_dim);
+ }
+ } else { // 2D columnwise
+ if (same_both_dims) {
+ group_dequantize_blockwise_2d_cw_kernel
+ <<>>(
+ input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors,
+ common_first_dim_blocks, K, total_row_blocks, first_logical_dim);
+ } else {
+ group_dequantize_blockwise_2d_cw_kernel
+ <<>>(
+ input_dptr, output_dptr, scale_inv_dptr, tensor_offsets_ptr, num_tensors,
+ common_first_dim_blocks, K, total_row_blocks, first_logical_dim);
+ }
+ }); // NOLINT(*)
+ ); // NOLINT(*)
+ NVTE_CHECK_CUDA(cudaGetLastError());
+}
+
+} // namespace fp8_blockwise
+} // namespace dispatch
+} // namespace transformer_engine
+
+#endif // TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_FP8_BLOCKWISE_CUH_
diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
new file mode 100644
index 0000000000..1fd1738f93
--- /dev/null
+++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
@@ -0,0 +1,1022 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+/*! \file group_quantize_fp8_blockwise.cuh
+ * \brief CUDA kernels to quantize grouped tensors with FP8 1D and 2D
+ * block scaling. A single launch walks 128x128 tiles across every tensor
+ * in the group, with each CTA decoding its owning tensor from the device-side
+ * GroupedTensor metadata. Supports SAME_BOTH_DIMS and VARYING_FIRST_DIM.
+ */
+
+#ifndef TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_BLOCKWISE_CUH_
+#define TRANSFORMER_ENGINE_GROUP_QUANTIZE_FP8_BLOCKWISE_CUH_
+
+#include
+#include
+#include
+#include
+
+#include
+
+#include "../../common.h"
+#include "../../recipe/recipe_common.cuh"
+#include "../../transpose/cast_transpose.h"
+#include "../../util/cuda_runtime.h"
+#include "../../util/ptx.cuh"
+#include "../../utils.cuh"
+#include "../core/common.cuh"
+
+namespace transformer_engine {
+namespace dispatch {
+namespace fp8_blockwise {
+
+using transformer_engine::detail::FP8BlockwiseColumnwiseOption;
+using transformer_engine::detail::FP8BlockwiseRowwiseOption;
+
+constexpr int kTileDim = 128;
+constexpr int kThreadsPerWarp = 32;
+constexpr int kThreadsPerBlock = 256;
+constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp;
+
+// ---- Per-expert scale layout helpers --------------------------------------------
+//
+// cuBLAS grouped FP8 block-scaling GEMM expects each expert's scales to live
+// in a contiguous per-expert sub-block of the global scale buffer:
+//
+// 1D rowwise CW (op_rowwise=true) per expert: (blocks_X, roundup(M_t, 4)) floats
+// 1D columnwise (op_rowwise=false) per expert: (blocks_y_t, roundup(K, 4)) floats
+// 2D rowwise (op_rowwise=true) per expert: (blocks_y_t, roundup(blocks_X, 4))
+// 2D columnwise (op_rowwise=false) per expert: (blocks_X, roundup(blocks_y_t, 4))
+//
+// The grouped kernel writes the GLOBAL buffer in tile-stride order, but the
+// position assigned to each tile must map into the per-expert contiguous
+// sub-block. These helpers compute the per-expert cumulative byte/float
+// offset and the per-expert local stride. SAME_BOTH_DIMS lets us derive
+// without walking offsets; VARYING_FIRST_DIM walks `tensor_offsets_ptr` once
+// per block (only the writing thread).
+
+__device__ __host__ constexpr size_t kScaleColAlign = 4;
+
+// 2D columnwise per-expert scale offset. Per-expert layout is
+// (blocks_X, roundup(blocks_y_t, 4)). Two paths, picked at call sites:
+// - SAME_BOTH_DIMS: direct formula (no walk, no reduction).
+// - VARYING_FIRST_DIM: CTA-cooperative prefix sum. Each thread accumulates a
+// partial over a strided subset of the tensors-before-this-one, a
+// warp-shuffle reduces inside each warp, then all threads sum the per-warp
+// partials read from a kNumWarps-element smem buffer to obtain the same
+// total. The non-linear DIVUP_TO_MULTIPLE on each per-tensor blocks_y_t
+// prevents a closed form.
+// ALL threads of the CTA must call this in lock-step (the cooperative path
+// contains __syncthreads()).
+template
+__device__ __forceinline__ size_t compute_2d_cw_expert_offset(
+ const size_t tensor_id, const size_t blocks_X, const size_t common_first_dim_blocks,
+ const size_t tile_row_stride, const int64_t* __restrict__ tensor_offsets_ptr,
+ size_t* warp_partials_smem, const int tid, const int warp_id, const int lane) {
+ if constexpr (kSameBothDims) {
+ return tensor_id * blocks_X * DIVUP_TO_MULTIPLE(common_first_dim_blocks, kScaleColAlign);
+ }
+ size_t my_partial = 0;
+ for (size_t i = static_cast(tid); i < tensor_id;
+ i += static_cast(kThreadsPerBlock)) {
+ const size_t blocks_y_i =
+ static_cast(tensor_offsets_ptr[i + 1] - tensor_offsets_ptr[i]) / tile_row_stride;
+ my_partial += DIVUP_TO_MULTIPLE(blocks_y_i, kScaleColAlign);
+ }
+ my_partial = warp_allreduce_sum(my_partial);
+ if (lane == 0) warp_partials_smem[warp_id] = my_partial;
+ __syncthreads();
+ size_t total = 0;
+#pragma unroll
+ for (int w = 0; w < kNumWarps; ++w) {
+ total += warp_partials_smem[w];
+ }
+ return total * blocks_X;
+}
+
+// 1D rowwise: per-expert layout (blocks_X, roundup(M_t, 4)).
+template
+__device__ __forceinline__ size_t expert_scale_offset_1d_rowwise(
+ size_t tensor_id, size_t blocks_X, size_t common_first_dim_blocks, size_t tile_row_stride,
+ const int64_t* __restrict__ tensor_offsets_ptr) {
+ if constexpr (kSameBothDims) {
+ const size_t M = common_first_dim_blocks * kTileDim;
+ return tensor_id * blocks_X * DIVUP_TO_MULTIPLE(M, kScaleColAlign);
+ } else {
+ // Each M_i is enforced to be a multiple of kTileDim (=128), hence a
+ // multiple of kScaleColAlign (=4), so DIVUP_TO_MULTIPLE(M_i, 4) == M_i and
+ // sum_{i(tensor_offsets_ptr[tensor_id]) / K;
+ return blocks_X * total_M_before;
+ }
+}
+
+// ---- Tensor-lookup helpers ----------------------------------------------------
+
+// Map a global tile-row index to its owning tensor. Delegates to the shared
+// `common::get_current_tensor_id` helper from `cast/core/common.cuh`. The
+// helper is parameterized by total `first_logical_dim` rather than per-tensor
+// block count, so we reconstruct it here for the SAME_BOTH_DIMS specialization
+// (VARYING_FIRST_DIM ignores it and uses `current_offset` + `offsets_ptr`).
+template
+__device__ __forceinline__ size_t find_tensor_id_by_block_y(
+ const size_t block_y_global, const size_t num_tensors, const size_t common_first_dim_blocks,
+ const size_t tile_row_stride, const int64_t* __restrict__ tensor_offsets_ptr) {
+ constexpr auto shape_rep =
+ kSameBothDims ? ShapeRepresentation::SAME_BOTH_DIMS : ShapeRepresentation::VARYING_FIRST_DIM;
+ const size_t first_logical_dim = num_tensors * common_first_dim_blocks * kTileDim;
+ const size_t tensor_id = common::get_current_tensor_id(
+ num_tensors, block_y_global * tile_row_stride, block_y_global, first_logical_dim,
+ /*last_logical_dim=*/0, tensor_offsets_ptr);
+ if constexpr (!kSameBothDims) {
+ // tensor_offsets_ptr carries cumulative element counts; tile_row_stride =
+ // kTileDim * K, so the per-tensor element span is divisible by
+ // tile_row_stride iff first_dim is a multiple of kTileDim.
+ if (tensor_id < num_tensors) {
+ const size_t span =
+ static_cast(tensor_offsets_ptr[tensor_id + 1] - tensor_offsets_ptr[tensor_id]);
+ if (span % tile_row_stride != 0) {
+ NVTE_DEVICE_ERROR(
+ "Grouped FP8 block-scaling quantize: each tensor's first dimension must be a "
+ "multiple of 128 (VARYING_FIRST_DIM).");
+ }
+ }
+ }
+ return tensor_id;
+}
+
+// Per-tensor block-y base for VARYING_FIRST_DIM (in 128-row block units).
+__device__ __forceinline__ size_t tensor_block_y_base_from_offsets(
+ const size_t tensor_id, const int64_t* __restrict__ tensor_offsets_ptr,
+ const size_t tile_row_stride) {
+ return static_cast(tensor_offsets_ptr[tensor_id]) / tile_row_stride;
+}
+
+// Per-vector amax. Uses bf16x2 `max.xorsign.abs` on sm_89+; FP32 fallback otherwise.
+template
+__device__ __forceinline__ CType compute_row_amax(const Vec& v) {
+#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+ if constexpr (std::is_same_v) {
+ static_assert(kVec % 2 == 0, "kVec must be even for packed bf16x2 amax");
+ const ptx::bf16x2* pairs = reinterpret_cast(&v.data.elt[0]);
+ ptx::bf16x2 amax_x2{static_cast(0.f), static_cast(0.f)};
+#pragma unroll
+ for (int p = 0; p < kVec / 2; ++p) {
+ ptx::abs_max_2x(amax_x2, amax_x2, pairs[p]);
+ }
+ return static_cast(__hmax(__habs(amax_x2.x), __habs(amax_x2.y)));
+ }
+#endif
+ CType amax = 0.f;
+#pragma unroll
+ for (int e = 0; e < kVec; ++e) {
+ amax = fmaxf(amax, fabsf(static_cast(v.data.elt[e])));
+ }
+ return amax;
+}
+
+// Per-tile column sum of the high-precision input -> one fp32 row at
+// dbias_workspace[tile_y_global * K + col]. 2 threads/column sum 64 rows each (combined via
+// shfl_xor); grouped_reduce_dbias later sums each expert's row-blocks. Tiles are always full
+// (experts are 128-row aligned), so all 128 rows are summed.
+template
+__device__ __forceinline__ void write_tile_dbias_partial(const IType smem_tile[][kTileDim],
+ const int tid,
+ const size_t global_col_base,
+ const size_t K, const size_t tile_y_global,
+ float* __restrict__ dbias_workspace) {
+ constexpr int kThreadsPerColDB = 2;
+ constexpr int kRowsPerThreadDB = kTileDim / kThreadsPerColDB; // 64
+ const int col_local = tid / kThreadsPerColDB; // 0..127
+ const int sub = tid % kThreadsPerColDB;
+ const int row_start = sub * kRowsPerThreadDB;
+ float partial = 0.f;
+#pragma unroll
+ for (int e = 0; e < kRowsPerThreadDB; ++e) {
+ partial += static_cast(smem_tile[row_start + e][col_local]);
+ }
+ partial += __shfl_xor_sync(0xffffffff, partial, 1);
+ const size_t c_global = global_col_base + col_local;
+ if (sub == 0 && c_global < K) {
+ dbias_workspace[tile_y_global * K + c_global] = partial;
+ }
+}
+
+// Per-vector multiply-and-quantize via fp32 intermediates.
+template
+__device__ __forceinline__ void quantize_row_vec(Vec