[PyTorch] torch.compile support for UnfusedDotProductAttention - #19
Open
pggPL wants to merge 34 commits into
Open
[PyTorch] torch.compile support for UnfusedDotProductAttention#19pggPL wants to merge 34 commits into
pggPL wants to merge 34 commits into
Conversation
…le + CUDA graphs
Refactor TE custom kernels used by the unfused attention path so that
`torch.compile(fullgraph=True, mode="reduce-overhead")` can trace the
forward and backward and capture them into CUDA graphs without graph
breaks.
- softmax.py / softmax.cpp: register all `scaled_*_softmax_{forward,backward}`
kernels as `torch.library.custom_op`s with fake impls and an autograd
binding that mirrors the previous `torch.autograd.Function`s. The C++
backward kernels now allocate a fresh output buffer instead of writing
in-place into `output_grad`, so the ops no longer alias their inputs
(required by `torch.library.custom_op` and inductor cudagraph trees).
- utils.py: convert `ConvertTHDtoBSHD` / `ConvertBSHDtoTHD` to
`torch.library.custom_op`s, with thin wrapper classes that keep the
existing `.apply(...)` callsite syntax. Drop the
`int(cu_seqlens[-1].item())` from the hot path of `ConvertBSHDtoTHD.apply`
-- under `torch.compile` it created an unbacked SymInt, which made the
Inductor partitioner emit `None` placeholders for output buffers and
caused `cudagraph_trees` to assert. `num_tokens` is now passed in by
the caller as a regular (Sym)Int.
- backends.py: in the THD branch of unfused DPA, capture
`total_tokens_q = query_layer.shape[0]` before overwriting
`query_layer` with the BSHD form, and thread it back into
`ConvertBSHDtoTHD.apply` at the end of the forward.
- test_torch_compile.py: add `test_unfused_dpa_torch_compile`,
parametrized over qkv layouts (`bshd_bshd_bshd`, `sbhd_sbhd_sbhd`,
`thd_thd_thd`, `bs3hd`, `sbh3d`), that compiles
`UnfusedDotProductAttention.forward` directly with `fullgraph=True,
mode="reduce-overhead"` and runs forward+backward several times so the
CUDA graphs are recorded and replayed.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor…to unfused_attention_torch_compile Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> # Conflicts: # tests/pytorch/test_torch_compile.py
…DotProductAttention
Make the FP8-emulation path (NVTE_UnfusedDPA_Emulate_FP8=1) of
UnfusedDotProductAttention traceable by torch.compile(fullgraph=True).
- backends.py: register the quantize+dequantize roundtrips used by
FP8EmulationFunc as torch.library custom ops
(te_fp8_emu::roundtrip_<QuantizerClass> and
te_fp8_emu::roundtrip_qkv_<QuantizerClass>) taking the quantizer as a
value-opaque argument, with fake impls for tracing. Ops are registered
only for the value-opaque quantizer classes
(Float8CurrentScalingQuantizer, MXFP8Quantizer); Float8Quantizer
(delayed scaling) carries scale/amax tensor state, is not
value-opaque, and deliberately keeps the plain eager path -- FP8
emulation with delayed scaling is not supported under torch.compile.
- backends.py: dispatch helpers `_fp8_emu_roundtrip{,_qkv}` key on
`type(quantizer).__qualname__` so they stay traceable for opaque
quantizer arguments; FP8EmulationFunc forward/backward now call them
(onnx_forward unchanged).
- backends.py: the joint q/k/v roundtrip clones any output whose
storage is shared with an input or another output, checking storage
identity directly -- the dequantized q/k/v can be views into one
combined buffer, and view metadata (`_base`) is not populated under
the torch-dispatch mode AOTAutograd runs custom ops with, so a
`_base`-guarded clone triggered the custom-op aliasing deprecation
warning under torch.compile.
- UnfusedDotProductAttention.forward: only query
FP8GlobalStateManager.get_fp8_recipe() when
fp8_meta["local_recipes"] is absent.
- test_torch_compile.py: add test_unfused_dpa_fp8_emulation_torch_compile
(current scaling + mxfp8, sbhd/bshd layouts; compiled fullgraph
forward+backward must match eager) and
test_unfused_dpa_fp8_emulation_delayed_scaling_eager guarding the
eager delayed-scaling path after the refactor.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>…output=True With fp8_output=True the backend returns a Float8Tensor -- a tensor subclass that cannot cross a torch.compile graph boundary -- so the forward dispatches to a torch._dynamo.disable'd wrapper, the same mechanism DotProductAttention and FusedAttention use module-wide. With fp8_output=False the dispatcher is resolved at trace time and adds no graph break. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…cudagraphs) Parametrize test_unfused_dpa_fp8_emulation_torch_compile over compile mode (default, reduce-overhead), run 3 iterations so the CUDA graphs are recorded and replayed. The te_fp8_emu roundtrip ops for current scaling are pure (no mutated args), so inductor cudagraphs capture them; verified no cudagraph skips with TORCH_LOGS=cudagraphs. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…n; run FP8 as an eager island FP8 in the unfused backend (emulation and Float8Tensor output) is not supported under torch.compile: the forward dispatcher routes fp8=True and/or fp8_output=True to a torch._dynamo.disable'd wrapper, same as DotProductAttention does module-wide. Remove the FP8-emulation compile tests. The te_fp8_emu::* custom ops taking value-opaque quantizers stay as the eager implementation of FP8EmulationFunc. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…Func The ops existed solely to make the FP8-emulation path traceable by torch.compile; since FP8 in the unfused backend now always runs as an eager island, they are dead machinery (plus import-time registration and output clones the plain eager path never needed). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The tex softmax kernels take 'float scale_factor' directly. The 0-D tensor wrapping was a leftover of the old autograd.Function idiom, where the float had to be a tensor only to fit save_for_backward; the custom ops keep the scale on ctx as a plain attribute. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…the callsite) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…cate dict, silence W0613 - run black over the four changed files (earlier commits skipped pre-commit) - drop unused 'import os' in test_torch_compile.py - drop duplicated module-level _default_causal_mask dict in softmax.py - del unused 'output' arg in the conversion setup_context helpers Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
for more information, see https://pre-commit.ci
* [JAX] Resync onto upstream PR NVIDIA#3036, restore TE-EP-only MoE block Reset 33 local commits onto phuong/ep-3-jax @ c34771d (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF ) Signed-off-by: Teddy Do <tdophung@nvidia.com> * tests/jax: trim TE-EP MoE suite (drop bootstrap, flax-wrapper, bias-zero) * drop ``TestZZZTeEpMoeBootstrap``: the re-bootstrap mismatch is a one-line guard in ``ep_bootstrap`` and not the MoE block's concern; exercising it from this suite also taints the per-process NCCL bootstrap cache for the rest of the file with no real upside. * drop ``TestTeEpMoEBlockFlax::test_init_apply_parity``: every config in ``_CONFIGS`` already runs ``MoEBlock`` (the Flax wrapper) end-to-end via ``test_forward`` / ``test_backward``, so this was a duplicate of ``softmax`` parity in another wrapper -- leave wrapper refactors to devs without paying for an extra CI run each time. * drop ``sigmoid-bias-zero``: with a zero-init bias buffer the routing math collapses to the no-bias case, so ``sigmoid`` already covers that numerical path. The bias-aware codepath is still exercised by ``sigmoid-bias-strong`` (non-zero bias). * refresh the module-level docstring to list intentional non-coverage so future readers don't re-add these tests. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/router: fix two bwd custom_partitioning bugs (aux-loss rank, topk closure) Two unrelated one-line bugs in the bwd custom_partitioning machinery that only surface once the MoE block's aux-loss path is lifted out of shard_map (the custom_partitioning_sharding_rule check is skipped under shard_map, which is why these never tripped before). 1. FusedMoEAuxLossBwdPrimitive.shardy_sharding_rule: ``grad_aux_loss`` is the cotangent of a scalar loss and is rank-0; declaring it with a spurious ``grad_one`` factor gave it rank-1 and tripped JAX's custom_partitioning_sharding_rule rank check at global view. Change the rule's third operand entry to empty: "const_buf_one, num_experts, grad_one -> i num_experts" -> "const_buf_one, num_experts, -> i num_experts" 2. FusedTopkWithScoreFunctionBwdPrimitive.partition: ``del result_infos, routing_map_format`` removed ``routing_map_format`` from the enclosing scope before the nested ``sharded_impl`` closure was invoked. Python closures resolve names at call time, not definition time, so when XLA finally invoked ``sharded_impl`` for the bwd partitioned impl it raised ``NameError: cannot access free variable 'routing_map_format'``. Drop ``routing_map_format`` from the ``del`` and leave a NOTE so future cleanups don't reintroduce the bug. Sibling partition methods (fwd topk, both aux-loss directions) already only ``del result_infos`` and need no change. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/ep: skip size-1 dp/fsdp axis in _ep_outer_axis A dp_resource or fsdp_resource that exists in the active mesh resource config but is sized 1 in the actual mesh would still be returned by ``_ep_outer_axis()``, pinning EP-output PartitionSpecs to a degenerate axis. JAX collapses size-1 mesh axes during lowering, which made the EP-output specs reference an axis that no longer exists at runtime -- breaking shard_map output stitching on configs where DP or FSDP is optional. Treat a size-1 axis as absent: prefer dp -> fsdp, but only when the candidate axis is actually sized > 1 in the current mesh. Falls back to the previous behaviour when no axis is configured at all. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/flax: realign _MoEBlock with post-resync moe() signature After the upstream PR NVIDIA#3036 resync the moe() API surface lost PermutationBackend (TE-EP is the only backend now), gate_inside_vjp (always True), and the per-call quantizer_sets knob (quantization flows through the standard TE autocast / with_quantizer_set context). It also gained apply_topk_weights_early and renamed the wrapper's private _align_size to the public align_size the test suite already uses. The Flax _MoEBlock wrapper was still passing the old kwargs, which broke every test that touched the wrapper. Wrapper changes: * drop "from ..moe import PermutationBackend" plus the dataclass field, the isinstance(..., PermutationBackend) validation in __post_init__, and the pass-through to moe(). * drop "from ..quantize import noop_quantizer_set" and the quantizer_sets=(noop, noop, noop) pass-through. * drop gate_inside_vjp=True. * rename _align_size: int = 0 -> align_size: int = 0 (matches what tests/jax/test_te_ep_moe.py already passes). * add apply_topk_weights_early: bool = False and pass it through to moe(). * refresh class docstring: drop permutation_backend / _align_size / quantizer_sets descriptions, add apply_topk_weights_early / align_size, note that quantization currently flows only through fp8_autocast. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: plumb token_counts to grouped_gemm and zero 0-token wgrad slices Two correctness fixes for the TE-EP MoE custom_vjp that together let the bwd parity tests pass on 0-token-globally experts, and drop a workaround that is no longer needed. (1) Plumb per-expert padded token_counts into grouped_gemm group_sizes. NCCL EP HT dispatch lays out recv_tokens expert-major as [expert_0_padded | expert_1_padded | ... | overalloc_tail] where each per-expert block already includes the dispatch_output_per_expert_alignment zero-padding and only the trailing overalloc tail (slack between sum(token_counts) and the worst-case recv_pr) is unused. Previously _ffn_fwd_per_shard built a static local_group_sizes = jnp.full((num_local_experts,), slots_per_expert), which over-counted by the overalloc tail and forced cuBLAS to run the GEMM for every group including 0-token-routed experts. Pipe the real per-shard token_counts (1, num_local_experts) from ep_prepare through _moe_fwd_rule (added to ffn_in_specs/ffn_in_args with ep2_spec), into _ffn_fwd_per_shard as token_counts_local, and reshape into local_group_sizes for both grouped_quantize and grouped_gemm. cuBLAS now skips both 0-token experts and the trailing overalloc tail. Mirror the residual spec change on the bwd (local_group_sizes residual moves from P() to ep2_spec). (2) Per-group jnp.where zero-fill on wgrad outputs. cuBLAS grouped_gemm skips groups with size_g == 0 without zero-filling the corresponding out[g, :, :] slice (cublaslt_grouped_gemm.cu lines 2086/2096). For a shard hosting an expert that received zero tokens globally, d_wo / d_wi_combined for that expert is left uninit, which propagates NaN straight into the user's optimizer state. Add wgrad_group_active = (local_group_sizes > 0)[:, None, None] in _ffn_bwd_per_shard and apply via jnp.where on d_wo (right after the wo wgrad) and d_wi_combined (right after the fused wi_0+wi_1 wgrad). Mask shape is (num_local_experts, 1, 1) so cost is negligible. (3) Drop the lax.cond zero-init guard on r_tok in _moe_fwd_rule._body. Previously a jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like) wrapper around recv_tokens worked around tex.ep_dispatch_fwd leaving the recv buffer uninit on fully-empty-receiver ranks. With (1) in place, cuBLAS skips experts whose group_sizes == 0 and the per-row trailing tail of dispatched recv_tokens is unread by every downstream consumer (subsequent grouped_gemms read only sum(group_sizes) rows; ep_combine and ep_dispatch_bwd are handle_mem-aware). The only per-row consumer that would propagate the tail is grouped_dbias (per-row segment_sum), which only runs when has_bias=True, and that FFN bias path is currently gated upstream (cuBLAS grouped_gemm has no fused bias on Hopper yet; PR 3083 adds the pure-JAX bias add). With (2) handling the user-visible wgrad-NaN risk on 0-token experts, the lax.cond is now redundant. Replace with a NOTE pointing at the two follow-ups that would force its reintroduction: - a future caller that reads the full recv tile non-group-aware (e.g. an inspect probe), or - the FFN bias path landing, which would resurrect grouped_dbias. Also rewrite the _ffn_fwd_per_shard and _ffn_bwd_per_shard docstrings to spell out the per-row vs per-group uninit semantics so the next person debugging a NaN here has the invariants written down. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR NVIDIA#3116) Address jberchtold-nvidia's PR NVIDIA#3116 nit "rename use_bias -> use_ffn_bias and use_expert_bias -> use_expert_routing_bias". The two flags are siblings (they enable two different bias buffers) but the old names suggested ``use_bias`` was the general fallback, which wasn't the intent. The new names make the FFN-vs-routing distinction obvious from the call site. * transformer_engine/jax/flax/moe.py use_bias -> use_ffn_bias (dataclass field + branch in __call__ + docstring entry) use_expert_bias -> use_expert_routing_bias (same) * tests/jax/test_te_ep_moe.py _make_block(use_expert_bias=...) -> use_expert_routing_bias sigmoid-bias-strong config key updated _reference_kwargs_from_config now reads use_expert_routing_bias ``_MoEBlock`` is still the experimental underscore-prefixed alias (no public ``MoEBlock`` export yet), so the rename is API-safe. The pre-resync legacy tests (``test_moe_vjp.py``, ``test_multiprocess_moe_vjp.py``) are intentionally not updated -- they already reference removed APIs like ``PermutationBackend`` and need a separate post-resync cleanup pass. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: address PR NVIDIA#3116 review feedback (hardcode align + expand inline justifications) Responds to jberchtold-nvidia's PR NVIDIA#3116 review threads on ``transformer_engine/jax/moe.py``. All changes are confined to a single file because each review thread targets a localized region and splitting mid-file would risk reordering bugs. Per review thread: 1. "Why do we need _with_sharding_constraint_cast_bwd? I haven't seen something like this required for our other VJPs." -- Expand the helper's docstring to spell out exactly why MoE needs it: unlike LN+MLP, the MoE bwd composes a bf16 cotangent from ep_dispatch_bwd with an fp32 cotangent from fused_topk_with_score_function_bwd (which the fwd's logits_2d -> fp32 promotion forces). Without the cast, ``d_x`` surfaces at fp32 even when ``x`` is bf16, doubling activation grad bandwidth and breaking any downstream LN bwd that pins a bf16 layout. (Review thread "Why do we need this utility function?".) 2. "Why is this dtype casting required? I don't recall us needing it for the non-MoE LNMLP block." -- Expand the comment above the bwd activation fp32 promotion to explain the MoE-specific math: LN+MLP's silu sits behind a downstream LN that absorbs the bf16 rounding error, while MoE's silu sits on the *expert* side of routing -- the bf16 rounding rides directly into expert_outputs and is summed across topk experts by ep_combine. Bf16 silu alone drifts ~1% vs fp32 silu and compounds through wo->combine into the ~1.4% per-element parity gap we measured against the pure-JAX softmax reference. Mirroring the fwd's fp32 promotion in the bwd keeps silu' in lock-step with silu. (Review thread on "# Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply".) 3. "Do we have a use-case for user-specified alignments beyond 128 currently? ... it'd make sense to instead hardcode _ALIGN_SIZE = 128 as a constant at the top of the file for now to simplify this MoEBlock API. We can always expand the API to support a user-specified align size in the future." -- Implement the suggestion. Drop ``align_size`` from ``_moe_fwd_rule`` / ``_moe_bwd_rule`` / ``_moe`` / public ``moe()``; shift the ``custom_vjp`` ``nondiff_argnums`` from ``range(9, 27)`` -> ``range(9, 26)``; replace ``effective_align = max(int(align_size), 128)`` with the new module-level ``_ALIGN_SIZE = 128`` constant. Trim the ``moe()`` docstring accordingly. (Review thread on "natural_spe = num_ep * max_tokens_per_rank".) 4. "Which axis name inputs are physical mesh axes and why can be logical axes? ... No need to make any changes for now, I just want to assess which are which and then we can discuss if it makes sense to support logical on some/all or if some are required to be physical axes." -- Add an "Axis-name parameters" section to ``moe()``'s docstring listing which kwargs are physical mesh axes (``ep_axis``, ``data_parallelism_axes`` -- they index ``Mesh.shape`` directly to compute ``num_ep`` / ``dp_size`` and to construct the ``P((dp..., ep), None, None)`` for ``jax.lax.with_sharding_constraint``) vs logical axes (``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` -- resolved via the Flax logical-axis rules). Also document why ``ep_axis`` / ``data_parallelism_axes`` are intentionally non-logical: the EP comm-group construction (``dp_color = rank // ep_size``) and the bootstrap signature check both require concrete integer sizes. (Review thread on "batch_pspec_axis = (*data_parallelism_axes, ep_axis)".) 5. "Is this NaN filtering a debugging artifact or something we need in the final version?" -- Strengthen the inline comment above ``sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, ...)`` to explicitly call this out as a CORRECTNESS REQUIREMENT, not a debugging artifact: it covers the sigmoid+K>1 underflow path where top-K sigmoid scores all round to zero and the ``weights / (weights.sum + 1e-20)`` normalisation emits NaN. Observationally the filter is a no-op on the dense unit-test distributions, but it must stay in for sparse / production routing. (Review thread on "sparse_probs = jnp.where(jnp.isnan(sparse_probs), ...).") Not addressed in this commit (intentional): * Review thread on the ``align_size: int = 0`` placeholder in ``flax/moe.py`` ("Placeholder comment for me to fix this so align_size is inferred automatically based on the recipe and doesn't need to be specified by the user"). That's jberchtold's own follow-up. * Review thread on the explicit ``tree_flatten`` / ``tree_unflatten`` on ``_Ctx`` ("better to use the ``@flax_struct.dataclass``"). Deferred to a separate, testable commit because changing a ``custom_vjp`` residual's pytree registration touches subtle ordering / None-handling semantics that warrant their own bisect surface. * Review thread on ``use_bias`` / ``use_expert_bias`` renames -- handled in the immediately preceding commit ``jax/flax,tests: rename use_bias/use_expert_bias for symmetry``. * Review thread on the ``expert_bias`` fp32 init -- already resolved during the Phuong PR NVIDIA#3036 resync (the redundant ``jnp.float32`` second-dtype argument on ``self.param`` was dropped; ``expert_bias`` now lives at ``self.dtype``). Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: strip PR-response framing from comments; drop sparse_probs NaN sanitizer * Rewrite the inline justifications added in 078a7d80 so each one reads as standalone code documentation, not as a reply to a reviewer: drop "per PR NVIDIA#3116 review", "review feedback", "Renamed from ... per PR ..." and similar PR/thread references from moe.py, flax/moe.py, and tests/jax/test_te_ep_moe.py. Technical content (why the fp32 promotion is needed for the MoE silu+multiply, why _with_sharding_constraint_cast_bwd exists, physical-vs-logical axis split in moe() docstring, the 128 alignment rationale) is preserved and reframed to be useful to a reader who has no PR context. * Drop the jnp.where(jnp.isnan(sparse_probs), 0, sparse_probs) guard. Tracing fused_topk_with_score_function.cu shows the kernel divides by sum_scores + 1e-20, so finite non-negative sigmoid scores cannot produce NaN here; the filter was only defense against upstream NaNs, which would mask a real regression if anything ever did start producing them. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: drop fp32 island around silu+multiply (fwd, bwd, reference) The SwiGLU intermediate (activation inputs gate_proj_out/up_proj_out, silu+multiply, and activation output) was previously promoted to fp32 in _ffn_fwd_per_shard and again in _ffn_bwd_per_shard, then cast back to the wi/wo GEMM dtype. The promotion bought nothing: the activation inputs come out of the wi grouped_gemm in bf16, the activation output is consumed by the wo GEMM (or wo's quantizer for FP8/FP4) in the same dtype, and storing higher precision than either consumer is wasted bandwidth. * _ffn_fwd_per_shard: drop the .astype(jnp.float32) on gate_proj_out and up_proj_out and the trailing .astype(sorted_x.dtype). The multiply now stays in the wi GEMM output dtype end-to-end. * _ffn_bwd_per_shard: symmetric simplification. jax.vjp(act_fn, ...) runs at bf16, both d_intermediate * silu' and d_intermediate * up stay at bf16, no casts. silu' is now consistent with silu (both bf16) so the chain rule composes cleanly without the prior fp32 detour. * tests/jax/test_te_ep_moe.py::_pure_jax_moe_reference: drop the matching fp32 silu in the parity reference so the test compares bf16-vs-bf16. Parity tolerance was not loosened; expect the comparison to tighten now that both sides round silu identically. Also fix an inaccurate inline comment at the apply_topk_weights_early fwd branch: the bf16 requirement on expert_outputs is enforced by ep_bootstrap (which rejects max_token_dtype != bf16 and sizes the NCCL EP HT mega-buffer for 2-byte slots accordingly), not by a runtime assert in the combine FFI. Signed-off-by: Teddy Do <tdophung@nvidia.com> * remove useless comments Signed-off-by: Teddy Do <tdophung@nvidia.com> * tests/jax: remove legacy MoE VJP tests + launcher; point CI at TE-EP successor test_moe_vjp.py and test_multiprocess_moe_vjp.py both import PermutationBackend from transformer_engine.jax.moe -- an API that was removed during the Phuong PR NVIDIA#3036 resync. Both files have been dead-on-import ever since; the multiprocess launcher run_multiprocess_moe_vjp.sh only points at the dead test. test_te_ep_moe.py (the TE-EP-only custom_vjp suite) already covers everything the legacy files exercised that is still meaningful: fwd, bwd parity vs the pure-JAX reference, aux loss, both score functions, multi-process. The legacy parametrize axis (PermutationBackend.PURE_JAX vs TRITON) no longer exists. * Delete tests/jax/test_moe_vjp.py * Delete tests/jax/test_multiprocess_moe_vjp.py * Delete tests/jax/run_multiprocess_moe_vjp.sh * qa/L0_jax_distributed_unittest/test.sh: switch the MoE VJP distributed suite invocation from run_multiprocess_moe_vjp.sh / test_multiprocess_moe_vjp.py to run_te_ep_moe.sh / test_te_ep_moe.py. * tests/jax/conftest.py: docstring reference updated. * tests/jax/test_te_ep_moe.py: drop stale "successor to ..." aside and the "mirroring run_multiprocess_moe_vjp.sh" parenthetical. Net: -981 / +9. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate Per reviewer feedback (Jaberchtold on PR NVIDIA#3036): the manual tree_flatten / tree_unflatten on _Ctx duplicate exactly what @flax.struct.dataclass auto-generates, and the permutation dataclasses elsewhere in this module already use flax.struct. Switching to @flax.struct.dataclass: * Removes ~75 lines of mechanical tree_flatten / tree_unflatten that have to be kept in sync with the field list by hand. * Keeps cfg as the single static field via flax.struct.field(pytree_node=False), so the fwd -> bwd boundary behavior under jax.custom_vjp is unchanged. * Drops two now-unused imports (dataclasses.dataclass, jax.tree_util.register_pytree_node_class) and adds flax.struct. Field order and the (children, aux_data) split are byte-equivalent to the previous manual implementation, so the pytree treedef seen by jax.custom_vjp is identical. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract Mirrors the sparse_probs NaN-sanitizer removal in fe44697: we trust ep_dispatch_fwd's contract that recv_topk_weights does not contain NaN, and would rather see NaN propagate (catching a contract violation immediately) than silently sanitize it. The mask_bool dance itself stays: ctx.expert_outputs and grad_pre_combine still carry NaN at padded slots (ep_dispatch_fwd leaves uninit memory in recv_tokens, FFN and combine_bwd propagate it), and IEEE NaN * 0 = NaN means jnp.where is structurally needed to overwrite padded positions with literal zeros before the sum reduction. What changed: * Drop `recv_w_clean = jnp.where(jnp.isnan(...), 0, ...)` and thread ctx.recv_topk_weights directly into w / mask_bool. * Replace the NaN-defensive comment block with a shorter note that explains the structural reason the mask is still needed (NaN in expert_outputs / grad_pre_combine at padded slots), without claiming anything about recv_topk_weights. Addresses Greptile P1 by removing the asymmetry (fwd had no sanitizer, bwd did) -- chosen direction is "remove the bwd sanitizer", matching the project-wide stance of trusting kernel contracts rather than papering over violations. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: assert output dtype; tests cover d_x parity (dtype + values) Two related dtype-contract changes: 1. moe.py: one-line assert at the moe() return path that output.dtype == x.dtype. Cheap structural guard against any future bug that lets the public output drift wider than the user-supplied input dtype. 2. test_te_ep_moe.py: extend test_backward to also check d_x, the gradient propagated back to the previous layer in backprop. _grad_step now uses jax.grad(loss_fn, argnums=(0, 1)) and returns (grads_variables, grad_x); the reference path does the same so we can compare. d_x is checked for: * shape == x.shape * dtype == x.dtype (protects the _with_sharding_constraint_cast_bwd wrapper that casts the fp32-promoted gate path back to the primal dtype on bwd; a regression in that wrapper would silently double activation gradient bandwidth) * finiteness + non-zero * numerical parity vs the pure-JAX reference d_x Addresses jberchtold review comment on test_te_ep_moe.py:650 ("we also need to check the final propagated gradient that will be passed onto the next layer in backprop"). test_combined_loss_grads is adjusted to ``grads, _`` unpacking; it doesn't need d_x for its main+aux finiteness check. Signed-off-by: Teddy Do <tdophung@nvidia.com> * tests/jax/test_te_ep_moe: strip docstring to just "what this suite covers" Drops two paragraphs whose content was agent-flavoured PR-review notes rather than user-facing test docs: * The final "FP8 / MXFP8 deferred" paragraph that referenced an internal review artifact (``.pr3036-review/INTEGRATION_DESIGN.md``) not in the repo. * The "Intentional non-coverage" section that explained which tests deliberately do not exist (no Flax-wrapper smoke, no re-bootstrap-mismatch test) and why -- exactly the kind of defensive / forward-looking justification prose CLAUDE.md says to keep out of the codebase. The remaining docstring covers what readers actually need: how to launch the suite, what each test class exercises, and a short note on the parametrize-vs-class layout. Addresses jberchtold review comment on test_te_ep_moe.py:54. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: address TE EP alignment review feedback Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: fix early topk weighting padded-slot masking Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: remove unused EP mesh size Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: tighten TE EP recv capacity bound Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: simplify late TE EP weighting Signed-off-by: Teddy Do <tdophung@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * jax/moe: reduce padded-slot recv weight masking Signed-off-by: Teddy Do <tdophung@nvidia.com> --------- Signed-off-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [Common] NVRTC for fused softmax and normalization (Phase 0)
Move the fused-softmax and LayerNorm/RMSNorm kernels from build-time template
instantiation to runtime NVRTC compilation, with full coverage of the existing
kernel set so the NVRTC path is the default.
Fused softmax:
- RTC compile/launch path for scaled / scaled-masked / scaled-upper-triangular /
scaled-aligned-causal softmax, keyed by dtype, shape and mask/causal mode.
- NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX (default OFF) restores the static
template dispatch.
Normalization (LayerNorm + RMSNorm, forward + backward):
- Replace the static REGISTER_NORM_LAUNCHER template fanout with an NVRTC
registry that compiles the selected (norm type, direction, dtypes, hidden size,
CTA config) kernel on first use and caches it.
- NVTE_BUILD_LEGACY_STATIC_NORM (default OFF) restores the static launchers.
- NVRTC-safe kernel sources: kernel sources/headers avoid common.h under
__CUDACC_RTC__; add the dtype aliases and a minimal std::is_same/conditional_t
in the RTC build, and replace a zero-length padding array (a GNU extension nvcc
accepts but NVRTC rejects) with a no-padding union specialization.
KernelManager (util/rtc.{h,cpp}) gains occupancy / function-attribute /
cooperative-launch helpers needed by the norm launchers.
Validated on sm_89 (RTX 6000 Ada): full normalization operator suite 192/192,
softmax + NVRTC unit tests pass; libtransformer_engine.so shrinks ~72 MB -> ~65 MB.
On sm_100a the NVRTC norm forward kernel builds where the static instantiation
crashed the compiler.
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* Add fully qualified name to softmax kernels
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* Add static fallback option, fix softmax acc_t dtype
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* add missing license
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* greptile changes
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* fix formatting, .clang-format
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* import cleanup
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* Test more columns for softmax, mr changes
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
* Fix tests
Signed-off-by: Carlos Gomes <cgomes@nvidia.com>
---------
Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt>
Signed-off-by: Carlos Gomes <cgomes@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>* support scaled swiglu, scaled srelu and scaled clamp swiglu Signed-off-by: zhongboz <zhongboz@nvidia.com> * vectorized loading improvement Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix bug for backward kernel Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * optimize Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix unit test failure Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update tests/cpp/operator/test_scaled_activation.cu Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * refactor, resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * address review comment Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * adaptive cta to fix slow block reduce for scale grads 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 * refactor to have gated and unary activation in activation infra 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 * reuse scale grad kernel for non scale grad since it is faster anyway 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 --------- Signed-off-by: zhongboz <zhongboz@nvidia.com> Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: vthumbe1503 <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> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
…DNN SDPA fprop (NVIDIA#3186) * [Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop cuDNN >= 9.24 SDPA (unified engine) accepts cumulative sequence lengths directly (cu_seq_len_q/kv) and can scale ragged offsets stored in coarser units back to elements via a per-tensor ragged offset multiplier. Use both in the f16/bf16 forward to skip the two conversion kernels (cu_seqlens_to_actual_seqlens and cu_seqlens_padded_to_offsets) that previously ran before every varlen fprop: - Bind the user's int32 cu_seqlens buffers as CU_SEQ_LEN_Q/KV for the padding mask, and the token-unit cu_seqlens_padded buffers as ragged offsets for Q/K/V/O/Stats with elements-per-token multipliers. - Gate on cudnn >= 9.24 and !dropout (the FE rejects dropout together with generated stats on the unified engine; TE always generates stats). CU_SEQ_LEN inputs pin implementation selection to the unified engine. - Keep the true batch size on the direct path: cuDNN reads the user's [actual_b+1] buffers, so the quantized max_b graph batch would read out of bounds. Token-dim bucketing (max_t) is unaffected. - No conversion workspace is needed on the direct path. - Factor the layout-group -> multiplier mapping into RaggedOffsetMultipliers (utils.h), shared by the graph builder and the legacy conversion kernel so the two cannot drift. The kernel rewrite also removes a cross-thread read (offsets_v[tid] = offsets_k[cu_seqlens_id]) that raced for quantized-batch tail entries with interleaved layouts. - Backward is unchanged (no backend support yet). NVTE_FUSED_ATTN_DIRECT_SEQLENS=0 disables the new path (testing aid, to be removed before merging). Validated on H100 and Blackwell against cuDNN 9.25: test_dpa_softmax_thd 15/15 in both modes, and direct-vs-legacy fused outputs/grads match for all THD layouts (thd_thd_thd, t3hd, th3d, thd_t2hd, thd_th2d) x MHA/GQA x padding/padding_causal x pad_between_seqs {false,true}. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [Common] Pass cu_seqlens directly to cuDNN SDPA FP8/MXFP8 fprop Extend the direct-seqlens path to the FP8/MXFP8 forward: bind the user's int32 cu_seqlens buffers as CU_SEQ_LEN_Q/KV for the padding mask instead of converting them to per-batch lengths with the cu_seqlens_to_actual_seqlens kernel before every call. (Unlike the F16 path, the FP8 path has no THD/ragged support, so this is the only conversion kernel there.) FP8/MXFP8 on the unified engine requires cuDNN >= 9.25 and cuDNN frontend >= 1.26. The frontend is header-only, so its version is a compile-time property; the gate uses a constant-folded CUDNN_FRONTEND_VERSION check (all referenced symbols exist in 1.25, so no preprocessor guards are needed). Dropout with generated stats stays on the legacy path, same as F16. Backward is unchanged (no backend support yet). Validated against cuDNN 9.25 + frontend 1.26 (test_dpa_fp8_vs_f16, padding configs, direct path on with no fallback): 56 passed on H100 (delayed + current scaling), 168 passed on Blackwell (adds MXFP8); zero failures. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [Common] Harden direct-seqlens version gates Address review feedback and a version-mix bug found in testing: - Remove the NVTE_FUSED_ATTN_DIRECT_SEQLENS env override (unnecessary; the version gates fully determine the path). - Check the compile-time CUDNN_VERSION in addition to the runtime version. The cuDNN frontend gates cu_seq_len support on min(compile-time, runtime) version, so e.g. a binary built against 9.24 headers running on a 9.25 library must take the legacy path; a runtime-only check let it attempt the direct fp8 graph, which the frontend rejects ("No suitable implementation") with no fallback. - Add a (currently redundant) CUDNN_FRONTEND_VERSION >= 1.25 check to the f16 gate for symmetry with the fp8 gate. - Raise the fp8 frontend floor from 1.26 to 1.27: 1.26 suffices for this C++ API use, but 1.27 is the floor for the python FE API's fp8 cu_seq_len support (exposed post-1.26-cut), and a single version story per feature avoids a silent gap when TE moves to the python FE API. Smoke-tested on H100: f16 THD 15/15 (direct path, cuDNN 9.24), fp8 padding subset 56 passed via legacy on 9.24, and 56 passed via legacy on the 9.24-compile/9.25-runtime mix that previously failed 56/56. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [Common] Lower fp8 direct-seqlens frontend floor to 1.26 Per TE team discussion: 1.26 is all the C++ FE API needs for fp8 + cu_seqlens (the support surface made the 1.26 cut; SDPA_fp8_attributes has had the setters since 1.25). Keep a comment noting that the python FE API requires 1.27 (its sdpa_fp8 binding gained cu_seq_len_q/kv post-1.26-cut), so a future migration to the python FE API knows to raise the floor. Smoke-tested on H100: f16 THD 15/15 (direct, cuDNN 9.24), fp8 padding subset 56 passed via legacy on 9.24 and on the 9.24-compile/9.25-runtime mix. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [Common] Fix sm120 THD softmax-stats layout and allocation use_ragged_stats excludes sm120, but the forward Stats declaration used the weaker condition (is_ragged_q && cudnn >= 9.6). On sm120 with THD, fwd therefore declared the ragged-style [b][s][h] stats stride with a null ragged offset (i.e. dense token-major), while bwd read the stats tensor as dense [b][h][s] -- a fwd/bwd layout mismatch. It would also have let the direct-seqlens path set a ragged-offset multiplier on a null ragged offset, a frontend validation error. Use use_ragged_stats for the fwd declaration so fwd and bwd agree, and give the stats allocation the same sm120 exception Max already has: without it the buffer is [num_tokens_q, h, 1], undersized for the dense [b, h, s_q, 1] graph whenever num_tokens_q < b * s_q. Pre-existing issue, independent of the direct-seqlens work. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [Common] Rename use_direct_seqlens to use_cu_seqlens_directly Clearer name for the flag controlling whether cu_seqlens buffers are passed straight to cuDNN SDPA; comment wording updated to match. No functional change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [Common] Suppress fn_size lint on fused_attn_arbitrary_seqlen_fwd_impl The direct-seqlens additions push the function to 508 non-comment lines, over cpplint's 500 limit. Per TE team, refactoring this long-standing function is beyond the scope of this PR, so suppress with NOLINT for now. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> * [Common] Pin the UNIFIED implementation on the direct cu_seqlens path cu_seq_len (and the ragged offset multiplier) are unified-engine-only, so with those inputs attached AUTO can only ever resolve to UNIFIED anyway. Pinning changes only the failure mode: an unsupported config fails with the unified engine's specific error instead of auto-selection's generic "no suitable implementation". Ordinary graphs (no cu_seq_len attached) keep AUTO. Matches the cudnn-frontend cu_seq_len sample, which pins for the same reason. Smoke-tested on H100: f16 THD 15/15 via the pinned direct path (cuDNN 9.24); fp8 padding subsets 56 passed via legacy on 9.24 and on the 9.24-compile/9.25-runtime mix. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Emil Gilliam <egilliam@nvidia.com> --------- Signed-off-by: Emil Gilliam <egilliam@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
NVIDIA#3204) [PyTorch] Add per-version FlashAttention env vars (NVTE_FLASH_ATTN_V2/V3/V4) NVTE_FLASH_ATTN enables or disables the whole FlashAttention family, but the choice between FlashAttention 2, 3, and 4 is automatic (package presence and compute capability) with no user override. Some workloads need to pin the FlashAttention generation, e.g. RL training that must produce bitwise-identical logprobs to an inference engine running a specific FlashAttention version: different generations use different tile sizes and online-softmax accumulation orders, so mixed versions between training and inference break batch-invariant / train-inference parity guarantees. Add NVTE_FLASH_ATTN_V2, NVTE_FLASH_ATTN_V3, and NVTE_FLASH_ATTN_V4 (default 1) that disable a specific FlashAttention version even when it is installed, following the existing NVTE_FLASH_ATTN filter pattern. Behavior is unchanged when the variables are unset. Signed-off-by: wdykas <wdykas@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* fix grouped linear hang Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * make the same change in grouped mlp as well Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
…GroupedLinear and fused grouped MLP (NVIDIA#3161) * Add optional caller-provided output/grad-input buffers to GroupedLinear module and fusible ops Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route per-op kwargs through Sequential via module-keyed op_kwargs mapping Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Write fused grouped MLP MXFP8 output and dgrad directly into caller buffers, eliminating the D2D copy + cleanup Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use 256-aligned splits in caller-buffer grouped MLP test Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * use basic_ops to track op kwargs Signed-off-by: YangFei1990 <feiw@nvidia.com> * add doc and resolve comments Signed-off-by: YangFei1990 <feiw@nvidia.com> * move out/dgrad_out out from the non_tensor_args Signed-off-by: YangFei1990 <feiw@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com>
* Fix FusedAdam empty tensor handling Signed-off-by: Jingyue Wu <wujingyue@gmail.com> * Move empty tensor filtering into MultiTensorApply Signed-off-by: Jingyue Wu <wujingyue@gmail.com> --------- Signed-off-by: Jingyue Wu <wujingyue@gmail.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
…opk_weight tensor (NVIDIA#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>
…3222) * Migrate NCCL EP submodule to NVIDIA/nccl-extensions * Drop PYTHONPATH override from EP test, example, and bench launchers * Drop cross-mode recv comparison in EP zero-copy IdentityAllSymm test Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * [Common] Rename 3rdparty/nccl submodule directory to nccl-extensions Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…NVIDIA#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>
update nccl-ext submodule name Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* Generalized Tensor Parallelism (GTP) init commit Co-authored-by: Jieming Zhang <jiemingz@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * GTP + gmm fusion Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * [fix] Respect per-op activation-offload markers in fused grouped MLP Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Code clean: rename GTP weight-sharding axis to gtp_remat Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert "[fix] Respect per-op activation-offload markers in fused grouped MLP" This reverts commit 8bb26f0. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Make TE GTP-agnostic at construction Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * GTP+nvfp4: fix GTP backward GEMM scaling-mode mismatch for bf16-gathered weights Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Make TE runtime GTP-agnostic via a DistributedWeight protocol Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Code clean - Take a single leader weight in the DistributedWeight dispatchers - Gather the FC2 grouped weight late in the fused grouped MLP Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Simplify the NVFP4 gather post-process; Materialize the EGTP FC1 weight before the NVFP4 dgrad dispatch Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Code clean - Rename gather coalescing flag grouped -> external_coalescing; - Clean up DistributedWeight wiring in TE modules - Restructure _all_gather_nvfp4 Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * fix comments Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Support DistributedWeight in the fusible grouped-linear ops path - Add a self-contained dispatch test with a fake DistributedWeight implementer Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Unify distributed-weight wgrad finalize to return a graph-safe dummy - `finalize_weight_grads` now accepts a weight list or a bare leader, mirroring materialize_weight_for_backward; - Centralize the in-place / dummy / async-None finalize contract in DistributedWeight.finalize_group_grads and delegate the dispatcher docstring to it. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> --------- Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Co-authored-by: Jieming Zhang <jiemingz@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
fix nproc Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
… docstrings Greptile P2: the ConvertTHDtoBSHD/ConvertBSHDtoTHD class docstrings said callsites keep the .apply(...) syntax without reflecting the actual argument list. Spell out the apply() signatures so the required args (incl. num_tokens / max_seqlen) are explicit. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…Attention The thd output conversion (q_format=='thd') passes total_tokens_q to the new ConvertBSHDtoTHD custom op, but total_tokens_q was only assigned on the training 'thd' input branch, not the inference 'thd_2bshd' branch, so thd KV-cache inference raised UnboundLocalError. Capture total_tokens_q once right after q_format is known, before any layout conversion: for both 'thd' and 'thd_2bshd' the query enters in thd layout so query_layer.shape[0] is the total query token count (a backed SymInt, unlike cu_seqlens_q[-1].item() which would sync the GPU and break torch.compile + cudagraphs). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
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 freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
14 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
Make the
UnfusedDotProductAttentionbackend traceable bytorch.compile(fullgraph=True, mode="reduce-overhead"), so the forward and backward can be captured into CUDA graphs without graph breaks.Scope:
torch.library.custom_ops with fake impls and autograd bindings; remove an unbacked-SymInt.item()from the hot path ofConvertBSHDtoTHD.fp8=True(emulation) and/orfp8_output=True(Float8Tensor output, a tensor subclass that cannot cross a graph boundary) the backend runs as an eager island — the forward dispatches to atorch._dynamo.disable'd wrapper, the same mechanismDotProductAttentionandFusedAttentionuse module-wide. FP8 attention always involves delayed scaling regardless of the recipe: S and dP are produced inside the kernel, so their amax cannot be known before quantization and they use delayed-scaling quantizers even under Float8CurrentScaling (seeDPA.init_fp8_metadata) — and delayed scaling (Float8Quantizer, tensor scale/amax state) is not supported under torch.compile.Type of change
Changes
softmax.py/softmax.cpp:scaled_*_softmax_{forward,backward}as custom ops; C++ backward kernels allocate a fresh output buffer instead of writing in-place intooutput_grad(custom ops and cudagraph trees forbid input aliasing).utils.py:ConvertTHDtoBSHD/ConvertBSHDtoTHDas custom ops;num_tokenspassed by the caller instead ofcu_seqlens[-1].item().backends.py:UnfusedDotProductAttention.forwarddispatches FP8 calls to an eager (dynamo-disabled) wrapper; the non-FP8 path is traced with no graph breaks.tests/pytorch/test_torch_compile.py:test_unfused_dpa_torch_compile(5 qkv layouts, fullgraph + reduce-overhead, fwd+bwd captured into CUDA graphs and replayed).Checklist:
🤖 Generated with Claude Code