[TE] IFU release v2.18 - #722
Open
aris134 wants to merge 18 commits into
Open
Conversation
Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>
…opk_weight tensor (#3187) * expose user-provided weights * adding pool based symm allocation; remove the persistent buffer in EpBuffer * add zero copy tests Signed-off-by: YangFei1990 <feiw@nvidia.com> --------- Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
… (#3171) * [Common/PyTorch] Support power-of-2 scales in grouped FP8 block-scaling quantize The default Float8BlockScaling recipe constrains scales to powers of 2, so the fused grouped path must honor the flag to stay numerically consistent with the unfused path. Thread a runtime pow_2_scales argument through the grouped quantize kernels (the shared scale helper already implements the rounding) and drop the force_pow_2_scales rejections. Also add a quantization-config parameter to nvte_group_quantize_dbias, which previously had no way to receive force_pow_2_scales or amax_epsilon on the bgrad path. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Enable fused grouped FP8 block-scaling path in GroupedLinear module Admit Float8BlockQuantizer in the fused GroupedTensor path on Hopper. The existing usage flags already match the Hopper TN-only mapping and the grouped GEMM selects transposed columnwise storage for NN/NT layouts, so only the path predicate changes. The fused path is an explicit opt-in via NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM, so raise on Blackwell (SM100/SM110) instead of silently falling back; the fused path has no MXFP8-broadcast emulation. Extend the fused dbias path (tex.bgrad_group_quantize) to FP8 block scaling when dgrad is required (dbias is computed in the rowwise pass). Add fp8_block_scaling to the fused-path tests with a Hopper-only gate, assert the fused path engages via a group_quantize spy, and add a Blackwell error-path test. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Enable FP8 block-scaling in GroupedLinear fusible op Replace the blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state with a per-op supports_float8_block_scaling flag and opt in the GroupedLinear op. Mirror the module-path predicate and fused-bgrad changes; since the graph-safe flow is default-on here (no env-var opt-in), other architectures fall back to the split-quantize flow instead of raising. Force use_split_accumulator=True for FP8 block-scaling operands in general_grouped_gemm_for_grouped_tensor, matching non-grouped general_gemm: cuBLAS has no fast-accum FP8 block-scaling algorithm, so the ops-layer forward failed algo selection without it. Add fp8_block_scaling coverage to the ops GroupedLinear tests. The CUDA-graph-safe test skips it for now: the replayed wgrad for the last expert diverges between replays depending on process allocation history; under investigation. Graph capture remains covered by the module-path test. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Use persistent workspaces in grouped-tensor GEMM general_grouped_gemm_for_grouped_tensor allocated its setup workspace (the cuBLAS per-group pointer/dimension arrays) and its cuBLAS workspace with per-call torch.empty. Under make_graphed_callables the forward and backward graphs share one capture memory pool, and a per-call allocation's block returns to that pool as soon as the Python reference dies, so blocks alias across the two graphs and captured kernels from one graph overwrite the GEMM metadata the other graph reads at replay. Observed as allocation-history-dependent failures in the ops-layer GroupedLinear cuda-graph test: capture-time cublasLtMatmulAlgoGetHeuristic NOT_SUPPORTED errors and corrupted wgrad outputs. This is also the likely mechanism behind the FP8 block-scaling wgrad corruption under CUDA graphs previously observed on Hopper and attributed to cuBLAS. Cache the setup workspace per (device, group size) and reuse the cached per-device cuBLAS workspace from the non-grouped path; consecutive GEMMs reusing one workspace are ordered by the stream. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Fix grouped FP8 block-scaling CUDA-graph deadlock via per-role cuBLAS workspaces The grouped-tensor GEMM path shared one persistent cuBLAS workspace across all grouped matmuls. cuBLAS's grouped GEMM keeps a grid-synchronization flag in the first bytes of that workspace and zeros it (via a captured memset) before each matmul. When the dgrad and wgrad grouped matmuls of a GroupedLinear backward share one workspace inside a replayed CUDA graph, that flag is aliased between the two matmuls; on the second graph replay the second matmul's cooperative kernel deadlocks with cuBLAS 13.6 (and corrupts the last expert's wgrad on cuBLAS < 13.6). The two matmuls are strictly stream-ordered (single stream, all-DEFAULT graph edges, no programmatic dependent launch), so this is shared-workspace reuse, not concurrent co-scheduling. Give dgrad/forward (slot 0) and wgrad (slot 1) distinct persistent cuBLAS workspaces, dedicated to the grouped path. Each slot remains a single persistent allocation, so CUDA-graph capture safety is preserved. Also drop the cuBLAS-version gate that skipped the FP8 block-scaling GroupedLinear CUDA-graph test, so it now exercises the fix on all supported cuBLAS versions. Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: document split-accumulator override, fix stale dbias comment - general_grouped_gemm_for_grouped_tensor: expand the comment to state that the fused grouped FP8 block-scaling GEMM forces use_split_accumulator=True and intentionally overrides the caller-supplied value, consistent with the Float8BlockScaling recipe (which fixes it True for fprop/dgrad/wgrad). - Float8BlockScaling recipe docstring: document that FP8 block scaling always uses split accumulation and that the fused grouped GEMM path ignores any caller- or recipe-supplied use_split_accumulator value. - GroupedLinear ops backward: correct the stale "BF16/FP16 path" comment; that branch also handles quantized paths where bgrad fusion did not apply (e.g. FP8 block scaling without a dgrad pass). Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Revert fusible-ops FP8 block-scaling; scope PR to GroupedLinear module Restrict this PR to the GroupedLinear module fused-quantize path. Revert the fusible-ops FP8 block-scaling enablement -- the BasicOperation opt-in gate, the GroupedLinear op support, and the fusible-ops test coverage -- back to main. Enabling fusible-ops FP8 block-scaling for both grouped and non-grouped paths is deferred to a separate PR. The blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state is restored. The split-accumulator guard in general_grouped_gemm_for_grouped_tensor is retained: it is correct for the module's FP8 block-scaling grouped GEMM. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Isolate grouped wgrad cuBLAS workspace by NT layout, not out-discreteness _get_grouped_cublas_workspace slots were keyed on is_discrete_out as a proxy for "this is the wgrad GEMM", which only holds when wgrad writes a list of per-expert grads. With single_grouped_weight=True, wgrad writes a single grouped weight-grad (GroupedTensor out, not a list), so is_discrete_out is False and it collided with dgrad on slot 0 -- reintroducing the FP8 block-scaling grid-sync-flag aliasing deadlock/corruption under CUDA-graph replay. Key the slot on the wgrad layout (NT / transb) instead: fprop (TN) and dgrad (NN) share slot 0, wgrad (NT) is always isolated on slot 1. Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: isolate grouped cuBLAS workspace per layout; drop redundant test spy - _get_grouped_cublas_workspace now keys the persistent workspace on the grouped GEMM layout, so fprop (TN), dgrad (NN), and wgrad (NT) each get a distinct workspace. The previous NT-vs-rest scheme left fprop and dgrad sharing one workspace; those have also been reported to conflict under CUDA-graph replay. Documents that the deadlock is deterministic and present through cuBLAS 13.7. - Drop the group_quantize call-counting spy in test_grouped_linear_grouped_tensor_path_matches_legacy; fused-path engagement is covered by the graph-safe test. Signed-off-by: Alp Dener <adener@nvidia.com> * updated grouped GEMM workspace comment on stale TMA descriptor related deadlocks Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
[PyTorch] Enable FP8 block scaling across fusible ops Enables FP8 block scaling (FP8BS) in composable fusible-ops for both non-grouped and grouped linears. - Drop the blanket FP8BS NotImplementedError in ops/op.py; FP8BS now uses the generic recipe-state path like MXFP8. - Add `Float8BlockwiseQTensorStorage.view()` so quantized norm outputs reshape without dequantizing (mirrors `MXFP8TensorStorage.view()`). - Route grouped FP8BS through the Hopper graph-safe path (cuBLAS 13.4+). Fuse dbias in the rowwise pass when dgrad is required. - Exercise FP8BS across the fusible-ops test suite with 128-divisible sizes. Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
…s with pickles (#3245) Allow Flash Attention tests to use checkpoints with pickles Signed-off-by: Tim Moon <tmoon@nvidia.com>
…#3240) * test: cover lazy NCCL setup for Newton-Schulz Construct the cuSOLVERMp context before any other distributed collective so the distributed test exercises lazy ProcessGroupNCCL communicator initialization. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * fix: initialize borrowed NCCL comm for Newton-Schulz Materialize the ProcessGroupNCCL communicator before borrowing its raw handle, retain the process group for the context lifetime, and validate the communicator size and rank during native context creation. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> --------- Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
* Add JAX attention tutorials Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Defer attention tutorial imports Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Pass query heads through attention wrapper Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Clean up attention tutorial docs Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add JAX attention tutorial landing page Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Refine context-parallel attention tutorial Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Refine attention tutorial references Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Skip large attention tests on older GPUs Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…P tests (#3269) * [PyTorch] Honor requested FA padding in CP tests The CP test runner creates inter-sequence padding when the FlashAttention padding case is requested, but it unconditionally told DPA that FlashAttention THD inputs had no padding. That mismatch left CP backward padding uninitialized.\n\nDerive the explicit padding state from the same condition used to generate the inputs and reuse it for both reference and CP calls, preserving the sync-free CUDA-graph path for non-padding cases. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
The FA-version job reloads checkpoints generated by test_attention.py, so the legacy delayed-scaling FP8 metadata is trusted inside this test. Mirror the L0 opt-in across every L3 attention execution path to preserve the secure runtime default while preventing false version-matrix failures. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
…CC 10.x (#3056) * Refine the support for D=256 on Blackwell server type GPUs Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add deterministic tests for D=256 for sm10.x Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Test clean up Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add PyT side tests for D=256 cuDNN fused attn on SM100-110 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Refine the PyT D=256 tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix the filtering condition for bias type for D=256 on sm10x for cudnn fused attn Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Code clean up Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Switching test check logic around Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add THD xfail marked tests for testing Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix the jax test bias condition fo skipping for D=256 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * nit: Fix comment Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Move the D=256 changes so as to follow an order of increasing cuDNN versions Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add support for D=256 THD in TE Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add TE JAX CP test support to verify D=256 for AG and Ring Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add PyTorch D256 CP tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: clean up comments Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Correct the atgs passed to test dpa Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Simplify THD fused attention layout check Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Rename PyTorch D256 fused attention tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Run JAX fused attention before softmax Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Relabel THD D256 cuDNN guard to 9.25 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fold D256 PyTorch CP tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Rename D256 PyTorch attention tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Remove D256 PyTorch attention param ids Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Disable FA4 D256 SWA on SM100 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix FA4 D256 packed QKV handling Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use is_training in DPA test helper Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* optimize for gemm's conditional enablement Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * avoid code repeatition Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove redundant test Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * address review comment Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Add NLTK pin <3.10.1 in encoder requirements Signed-off-by: Alex Y. Chan <alechan@nvidia.com>
Merge NVIDIA TransformerEngine release_v2.18 (tag v2.18, VERSION 2.18.0)
into the ROCm IFU-2.18 landing, following the IFU-release_v2.17 pattern.
Conflicts resolved with ROCm adaptations (matching the in-progress 2.19 IFU):
- pytorch/module/{linear,layernorm_linear,layernorm_mlp}.py: keep the
IS_HIP_EXTENSION columnwise set_usage branch in _get_weight_quantizers;
drop the old preswizzle block (upstream relocated it into the
_enable_weight_preswizzle forward-path helper).
- common/cast/dispatch/quantize.cuh: keep the #ifndef __HIP_PLATFORM_AMD__
grouped block-scaling guard (NVTE_ERROR on ROCm); adopt upstream's
force_pow_2_scales call signature; drop the now-obsolete force_pow_2_scales
pre-check.
- tests/jax/test_fused_attn.py: keep not is_hip_extension() guard using
upstream's compute_capability local.
- tests/pytorch/attention/run_attention_with_cp.py: keep ROCm sys.path.append
shim plus upstream's multiline import.
- tests/pytorch/test_grouped_linear.py: keep IS_HIP_EXTENSION skip; take
upstream's updated comment.
Upstream v2.18 (#3232) relocated weight scale-swizzle enablement into the forward-path _enable_weight_preswizzle helper, called unguarded in Linear, LayerNormLinear, LayerNormMLP, and GroupedLinear. On ROCm the pre-merge code never set weight optimize_for_gemm=True (it lived only in the CUDA-only branch of _get_weight_quantizers); the new helper returns True for MXFP8, which would newly request with_gemm_swizzled_scales on ROCm quantizer paths that reject it. Guard the helper to return False on ROCm, preserving pre-merge behavior at all four call sites. ROCm continues to drive scale-swizzling per weight quantizer in _get_weight_quantizers via set_usage/keep_fp8_weight_transpose_cache.
aris134
marked this pull request as ready for review
September 3, 2026 14:34
aris134
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
September 3, 2026 14:34
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
11 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: