Ifu dev 20260803 v2.18 merge - #720
Merged
Merged
Conversation
Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
* feat(common): add dense topk index router output Signed-off-by: Harry Zhou <[email protected]> * [Common] Optimize dense fused router backward Signed-off-by: Harry Zhou <[email protected]> * [Common] Align dense router fallback with p3R Signed-off-by: Harry Zhou <[email protected]> * [PyTorch] Support int16 weak refs for CUDA graph reuse Signed-off-by: Harry Zhou <[email protected]> * [Common] Guard dense router topk index APIs Signed-off-by: Harry Zhou <[email protected]> * [Common] Harden dense router API guards Signed-off-by: Harry Zhou <[email protected]> * [PyTorch] Clarify dense router format guards Signed-off-by: Harry Zhou <[email protected]> * [PyTorch] Preserve router leading dimensions Signed-off-by: Harry Zhou <[email protected]> --------- Signed-off-by: Harry Zhou <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>
…g Last/Both Dims in Group Quantize (#3114) * Add grouped FP8 tensor-scaling quantization Route grouped Float8CurrentScalingQuantizer through the existing grouped quantize entry point, prepare per-group current-scaling metadata with existing amax/scale helpers, and add focused tests plus a GB200 bandwidth benchmark. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_5507e814ee50f9ff304a4ce708d19768 Orchestra-Run: run_516e1e26891f4ce7d4cde07147c10862 * Optimize grouped FP8 tensor-scaling quantization Use wider vectorized grouped FP8 cast-transpose tiles and vectorized masked stores for rowwise and columnwise outputs. Capture all benchmark modes in a single post-warmup profiler range. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_3d6e33eab11e293d72eb4394bad76a81 Orchestra-Run: run_a6e2c31d5fdf850594f71438e53148da * Fix grouped FP8 bias backward routing Route non-MXFP8 grouped-linear bias backward through group_quantize plus grouped dbias while keeping MXFP8 bgrad_group_quantize fusion intact. Add focused zero-row grouped FP8 coverage and a current-scaling GroupedLinear bias-backward regression. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_ab566800d87047635cd27f9e64661abe Orchestra-Run: run_5f9bfef17ccd854232c54d56268ef9e8 * Optimize grouped FP8 tensor-scaling cast kernel Use packed FP8 conversion and reduce columnwise transpose staging register and synchronization overhead in group_cast_fp8_kernel. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_7a830e018ceac8de0018280bd0740a54 Orchestra-Run: run_d2f1df4ffc2265d9cfa5ed01028ee476 * Fix grouped FP8 vector conversion deduction Match the grouped FP8 conversion helper's element-count template parameter to Vec's uint32_t parameter so rowwise, columnwise, and activation instantiations can build. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_30c4b6ddb896e5ea3ca5b54731d2c819 Orchestra-Run: run_e95cdbb445943304622b95736f0eca49 * Optimize grouped FP8 overallocated launch and transpose staging Use cached grouped offsets to avoid launching FP8 quantization over unused overallocated rows, permit larger grouped backing buffers when split metadata is present, and tighten full-tile vector paths in the grouped FP8 cast kernel. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_c5db93823dc101838cb1323e283cd6e9 Orchestra-Run: run_063e2e4c724e132612aa5597d6765c9b * Size grouped FP8 launches from active output shape Use the FP8 grouped output logical shape when computing the tensor-scaling launch grid so overallocated buffers with active metadata avoid empty tail-row launches while preserving the allocated-shape fallback. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_b4abb47c990404d73142342a19996a3f Orchestra-Run: run_8f09e7b9d7af9754ef505f2e2ce3cf90 * Improve grouped FP8 columnwise quantization bandwidth Use larger grouped FP8 tiles with 8-warp CTAs and 16-row columnwise store fragments. Treat uniform overallocated FP8 grouped outputs as same-shape wrappers during output reuse so the timed path avoids varying-shape metadata overlaunch. Add overallocated current-scaling coverage for all grouped FP8 direction modes. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_3f98ac9c5b82192ec289d8d2a9816c7f Orchestra-Run: run_83f3b99cc950024cf06ee836337fbf72 * Optimize grouped FP8 columnwise staging Stage columnwise transpose fragments through shared-memory vectors with smaller columnwise row tiles to reduce register pressure and barrier overhead while preserving the larger rowwise-only store path. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_495cc57eef84749103aded403a508d99 Orchestra-Run: run_53e038e90f83186bc6c12cb722c986b5 * Optimize grouped FP8 same-shape quantization Add fast grouped FP8 rowwise and full-tile columnwise paths for uniform active groups while preserving the general fallback for varying grouped metadata. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_4c33e88776c8a7148e9da5cc2bae84ea Orchestra-Run: run_2caaff219394eb5d59b7be38ab2bf346 * Optimize grouped FP8 bidirectional full-tile quantization Add a same-shape bidirectional full-tile kernel with wider input vectors and rowwise stores while preserving the existing rowwise-only, columnwise-only, and fallback grouped paths. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_87cec01d94f053b53e3c79377ad379ab Orchestra-Run: run_ed48db00a730a4bf56530d551ecd350e * Use compact full-tile grouped FP8 bidirectional quantize Route same-shape rowwise+columnwise grouped FP8 tensor-scaling quantization through the compact full-tile transpose schedule instead of the wide dynamic-shared-memory variant, preserving the existing single-direction and fallback paths. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_fdddd228a620039c024b4ecf43f3ab42 Orchestra-Run: run_30a2753eea9c893cb0fadb8233da8ce6 * Use streaming stores for bidirectional grouped FP8 rowwise output Hint the rowwise stores in the full-tile rowwise+columnwise grouped FP8 path as streaming global stores to reduce cache/writeback pressure without changing single-direction launch geometry. Orchestra-Work-Order: wo_aea2e337b06582111bba66a6d6158a9e Orchestra-Task: task_bf82020032e68276f4e47c65f62d97ae Orchestra-Run: run_754ea4c864f329c6f2003b413b723c43 * Handle varying-dim grouped FP8 quantization feedback Add graph-safe grouped FP8 tensor-scaling metadata, support varying last dimensions, preserve same-shape fast paths, adjust grouped FP8 columnwise allocation by architecture, and expand benchmark/test coverage for the reviewed shape cases. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_d104e74844fbc3d3b1a98a8d96d76037 Orchestra-Run: run_1314e997c61ffb92ff7120b0b26f0318 * Handle nonaligned varying-last grouped FP8 quantization Map varying-last columnwise tiles per group to avoid tile-alignment device errors, expand nonaligned boundary coverage, and restore same-shape benchmark baseline criteria. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_14e0e7973300d26f69550bc0aee21acc Orchestra-Run: run_2f42b8ba138ed8b2b4d9dc90b92caf85 * Emit same-session grouped FP8 baseline measurements Add grouped FP8 benchmark support for baseline-ref same-session reports and update the benchmark request to enforce same-shape baseline regression checks alongside the per-mode throughput thresholds. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_d0cada957a4aafdce9d52be86520e182 Orchestra-Run: run_4da74e9bdb4f4a4c72304a385692b6c9 * Build grouped FP8 benchmark baseline extension Update the grouped FP8 benchmark driver so same-session baseline checks out and builds the baseline ref into an isolated PyTorch install, verifies the baseline subprocess loads those shared objects, and preserves the required same-shape baseline comparisons. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_4fd88b172872f547f2f2d0053dce73d1 Orchestra-Run: run_6a44ee0467ffff47d4b278de6127354d * Fix grouped FP8 targeted correctness failures Preserve grouped delayed-FP8 amax metadata and keep unsupported FP8 tensor-scaling quantizers out of the grouped GEMM path. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_2aa8e6bf11ae356f4b34d4540b508031 Orchestra-Run: run_302681098d7f4e05b0ad96450f2d9826 * Make GroupedLinear single-param state-dict tests hermetic Set NVTE_GROUPED_LINEAR_SINGLE_PARAM inside the targeted state-dict tests so they exercise the gated single grouped parameter path without relying on external environment setup. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_261900f987bdc9397965019983a77c41 Orchestra-Run: run_c6624e34717cbe121b3e0edcf490e3d3 * Specialize varying-first grouped FP8 quantization Add a segmented flat rowwise kernel for varying-first grouped FP8 tensor-scaling outputs while preserving the existing same-shape fast path. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_c1b7020b27290318848ef6ac9048dd5f Orchestra-Run: run_5c257b8a5d2e7e4aa95e67aa16436166 * Fix grouped FP8 benchmark baseline compatibility Omit the last_dims keyword when absent so the same-session baseline can run against the base extension, and refresh the benchmark request to include direct varying-last current-scaling coverage. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_c20a3c94fdc798e741a469bd7bb9c4df Orchestra-Run: run_457448e6cba80fc63ac72b3db71c5fd0 * Optimize varying-first grouped FP8 quantization Dispatch varying-first tensor-scaling work per group to reduce inactive-tail CTAs and offset lookup overhead while preserving same-shape fast paths and graph-safe device metadata handling. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_d84e1fefef8641e558df064452f4689b Orchestra-Run: run_a361ca2f93fcec53ddd60dd99f4639e5 * Optimize aligned varying-first grouped FP8 quantize Add a no-tail rowwise flat kernel for aligned varying-first grouped FP8 tensor-scaling quantization and keep same-shape and varying-last dispatch isolated. Tighten benchmark profiler timing so post-warmup measured ranges exclude profiler start overhead. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_2e478be1fb38195f36d25c51320dc01f Orchestra-Run: run_9a133a75fa3d98dc3b1a63b0ff4d84af * Protect grouped FP8 benchmark raw report output Write grouped FP8 benchmark reports to a sidecar path by default and label script reports as benchmark_raw_report/v1 so regular 100-iteration measurements are fetched instead of the wrapper command report. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_27770b2e1d490b1a3053244d4b4ce248 Orchestra-Run: run_214052d0c1316e231443d645183a2675 * Mirror grouped FP8 benchmark report for Orchestra fetch Write the grouped FP8 benchmark JSON once and mirror the completed sidecar to ORCHESTRA_BENCHMARK_RAW_REPORT when running under Orchestra so the benchmark fetch path can parse the emitted measurements. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_b2e2747371204088c8e3f7cf10263164 Orchestra-Run: run_1d4ea38266807c8acb59143ee74ba241 * Write grouped FP8 benchmark report to raw path Allow the grouped FP8 benchmark to use ORCHESTRA_BENCHMARK_RAW_REPORT as its primary output so the benchmark wrapper can fetch canonical measurements directly. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_10fdcfef6b70de4676b7843e4bbfac31 Orchestra-Run: run_4ce57df9e86d6d03a26f7aa95ac252cc * Emit compact grouped FP8 raw measurements Write canonical grouped FP8 benchmark measurements to ORCHESTRA_BENCHMARK_RAW_REPORT in a small schema-shaped payload so the benchmark wrapper can materialize per-mode threshold evidence. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_3e862eebd585c74f2a58497fedea3511 Orchestra-Run: run_3770ab3dbbf51329d0839b3d10a91b5c * Preserve grouped FP8 raw benchmark results Write candidate_results and nonempty measurements into the Orchestra raw report path, and fail fast if the benchmark cannot produce threshold-ready evidence. Orchestra-Work-Order: wo_9d18259ce6c2833da1178606e08d251a Orchestra-Task: task_aa587a7b0d35aa9c2b715ec1b7c8bec3 Orchestra-Run: run_b42870e5d5e142a6cbf53bb5a3cafc2e * changes to improve amax kernel Signed-off-by: Varun Thumbe <[email protected]> * cleanups Signed-off-by: Varun Thumbe <[email protected]> * cleanup tests Signed-off-by: Varun Thumbe <[email protected]> * further cleanup Signed-off-by: Varun Thumbe <[email protected]> * further cleanup Signed-off-by: Varun Thumbe <[email protected]> * profile code push for now Signed-off-by: Varun Thumbe <[email protected]> * cleanups Signed-off-by: Varun Thumbe <[email protected]> * grouped amax in a seperate file Signed-off-by: Varun Thumbe <[email protected]> * changes so far Signed-off-by: Varun Thumbe <[email protected]> * clean Signed-off-by: Varun Thumbe <[email protected]> * no need for all same first Signed-off-by: Varun Thumbe <[email protected]> * dead code Signed-off-by: Varun Thumbe <[email protected]> * revert Signed-off-by: Varun Thumbe <[email protected]> * clean Signed-off-by: Varun Thumbe <[email protected]> * all changes: Signed-off-by: Varun Thumbe <[email protected]> * cleanup Signed-off-by: Varun Thumbe <[email protected]> * cleanpus Signed-off-by: Varun Thumbe <[email protected]> * Simplify splits_to_offsets and avoid modifying the multi-offset API - Reverted all changes to splits_to_offsets_multi and nvte_splits_to_offsets_multi. - Adapted the existing `kernel` in splits_to_offsets.cu using template parameters instead of introducing a separate kernel, preserving original kernel name and layout. - Kept the optimized nvte_splits_to_offsets behavior with varying first and last dims. * resolve merge conflicts Signed-off-by: Varun Thumbe <[email protected]> * add the single grouped weight skip check Signed-off-by: Varun Thumbe <[email protected]> * some optimizations.. unclean Signed-off-by: Varun Thumbe <[email protected]> * cleanups Signed-off-by: Varun Thumbe <[email protected]> * optimize amax kernel Signed-off-by: Varun Thumbe <[email protected]> * optimize cast kernel Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * improve colwise cast perf Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nvte tensors in the API Signed-off-by: Varun Thumbe <[email protected]> * address review comments Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments, fix lint errors Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * no need for type_trait Signed-off-by: Varun Thumbe <[email protected]> * varying last dims can also be overallocated Signed-off-by: Varun Thumbe <[email protected]> * split commmon into layout and tma files for arch specific vs non arch specific Signed-off-by: Varun Thumbe <[email protected]> * address reviem comment Signed-off-by: Varun Thumbe <[email protected]> * no need to depend on multi tensor impl.. nvte API to compute grouepd scale from amax Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test Signed-off-by: Varun Thumbe <[email protected]> * Remove duplicate comment in transformer_engine.h Removed duplicate brief comment about scaled prefix-sum offsets. Signed-off-by: vthumbe1503 <[email protected]> * fix comment and add nvte check Signed-off-by: Varun Thumbe <[email protected]> * add CPP tests for current scaling group quantize + splits to offsets 2d Signed-off-by: Varun Thumbe <[email protected]> * checks for tensor offsets as well Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Signed-off-by: Varun Thumbe <[email protected]> Co-authored-by: Orchestra <orchestra@local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Varun Thumbe <[email protected]>
…s (#3093) * [PyTorch] NVFP4: cache GEMM-swizzled weight scale factors across microbatches For block-scaled NVFP4 a cached weight participates in two GEMMs per step: fprop (rowwise scales) and dgrad (columnwise scales). The GEMM-ready scale swizzle was recomputed lazily inside every GEMM and discarded, so with N microbatches the weight scale swizzle ran 2*N times per step even though the weight is quantized only once. Because weight RHT is disabled, the weight scales are not swizzled by the cast-fusion path; with optimize_for_gemm off they also skip the post-quantize fallback swizzle, so the only swizzle site left for the weight is the lazy one inside general_gemm (swizzle_scales_for_gemm), which re-runs on every GEMM. (Activation input/grad_output quantizers already set optimize_for_gemm=True, so they were pre-swizzled via cast-fusion/fallback; only the weight was missed.) Set weight_quantizer.optimize_for_gemm=True on the cached, non-FSDP path so the swizzle is done once at quantize time (via the post-quantize fallback), persisted on the cached workspace (_with_gemm_swizzled_scales=True), and reused by every GEMM (swizzle_scales_for_gemm early-returns) -> 2 swizzles per step instead of 2*N. Applied to Linear, LayerNormLinear, LayerNormMLP (fc1+fc2) and GroupedLinear (per expert). Gated to the cached path (is_first_microbatch is not None) with fsdp_group is None and not is_fsdp2: FSDP/FSDP2 all-gather weights using the un-swizzled scale layout, so pre-swizzling is unsupported there. No-op for recipes whose scales do not require swizzling (e.g. per-tensor FP8). Swizzling is a pure layout permutation, so numerics are unchanged. Add tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py verifying the cached eager-swizzle path matches the lazy-swizzle baseline (fprop + dgrad) for Linear/LayerNormLinear/GroupedLinear and that the swizzled flag is persisted. Signed-off-by: Cael Ling <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review on NVFP4 cached-weight scale swizzle Drop the FSDP/FSDP2 gating on optimize_for_gemm in Linear, LayerNormLinear, LayerNormMLP and GroupedLinear. FSDP1 replicates (does not shard) the scale factors, so the swizzle layout is irrelevant there, and weights are not cached under FSDP2; the guard only added a misleading comment and dead conditions. Pre-swizzle the weight scales whenever the quantized weight is cached. Tests: - Fold the GroupedLinear case into the parametrized test_weight_swizzle_cache_numerics by passing m_splits only for GroupedLinear, removing the duplicated grouped-only test. - Add LayerNormMLP coverage (fc1 + fc2 two-quantizer path), generalizing the cached-workspace-count assertion per module type. - Parametrize test_lazy_path_not_swizzled over all four module kinds. Signed-off-by: Cael Ling <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * enable weight swizzling for most cases Signed-off-by: Varun Thumbe <[email protected]> * mxfp8 zero kernel for noop cuda graph compat Signed-off-by: Varun Thumbe <[email protected]> --------- Signed-off-by: Cael Ling <[email protected]> Signed-off-by: Varun Thumbe <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Varun Thumbe <[email protected]>
…llelism (#2829) * [PyTorch][CP] Fix THD AllGather CP: offset-based approach with proper cu_seqlens - Use per-step cu_seqlens_q_padded to select Q chunks instead of tensor slicing - Use padded cu_seqlens_kv for K/V reordering (ensures divisibility) - Add cu_seqlens_kv and cu_seqlens_kv_padded to AllGather function signature - Compute per-step Q and KV cu_seqlens correctly from actual seqlens - Support non-causal attention (all KV visible) - Zero-initialize out/dq for THD to avoid garbage in padding regions - Save per-step cu_seqlens in ctx for backward (avoid recomputation) Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch][CP] Enable THD+all_gather tests in test_attention_with_cp Remove skip gates that blocked THD format with all_gather CP comm type. Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch][Fused Attn] Fix max_logit masking for non-zero-starting cu_seqlens_q_padded The interleaved valid mask computation assumed cu_seqlens_q_padded starts at 0. With the CP offset-based approach, cu_seqlens_q_padded can start at a non-zero offset, causing a size mismatch. Use absolute positions from cu_seqlens_q_padded to build the valid mask instead. Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * some cleanup of ag+thd impl and gate e e te test for flash+ag+thd Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * improve the logic and remvoe for loop from the code Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * AG+THD SWA: extend KV visibility for right window and rename a2a-specific helpers The AllGather THD path was not extending KV visibility beyond the causal boundary when window_size had a right component > 0, meaning tokens right of the diagonal were invisible to the kernel. Fix by adding window_size[1] to visible_padded (clamped at actual seqlen) and max_seqlen_kv_. Also rename reorder helpers to backend-neutral names since AllGather now uses them too, and add a clarifying comment for non-causal KV cu_seqlens. Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Add pad_between_seqs support for FlashAttention 3 with CP Add support for padding between sequences (pad_between_seqs) in the FlashAttention 3 backend when used with context parallelism (CP). Key changes: - backends.py: Pass fa_pad_between_seqs through to FA3 forward/backward - context_parallel.py: Handle pad_between_seqs in A2A and P2P CP paths, zero FA3 padding garbage in CP forward, fix a2a backward alignment - dot_product_attention.py: Auto-detect pad_between_seqs from cu_seqlens - utils.py: Gate FA3 deterministic backward for hdim>=256, fix flash_attn_supported override for cross-attention and large head_dim, disable UnfusedDotProductAttention for pad_between_seqs, add SM100+ FA3 skip Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Add pad_between_seqs tests for CP and non-CP FlashAttention Add test parametrization for pad_between_seqs in flash attention tests. Update run_attention_with_cp.py to support the new parameter and fix batch boundary alignment in the non-CP FA3 path. Run tests in parallel when multiple GPUs are available. Signed-off-by: Sudhakar Singh <[email protected]> * [QA] Add CP deterministic tests to L3 and support TE_PATH in FA test Add deterministic CP test runs to L3 FA versions test. Support TE_PATH positional arg and fix GPU threshold for parallel test execution. Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Fix FA3 deterministic gate to match upstream backward constraint The previous check disabled FA3 for deterministic mode whenever head_dim_qk > 128, which was overly conservative — FA3 forward supports deterministic execution at any head dim. The actual constraint from flash_api.cpp is that the backward pass does not support deterministic mode when max(head_size, head_size_v) >= 256. Narrow the gate to only disable FA3 during training (backward) and raise the threshold to >= 256, checking both head_dim_qk and head_dim_v to handle MLA configs with asymmetric head dimensions. Ref: https://github.com/Dao-AILab/flash-attention/blob/ac6f2eb5/hopper/flash_api.cpp#L1370 Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Disable FlashAttention 4 for pad_between_seqs with THD The pad_between_seqs gate in get_attention_backend only disabled FlashAttention 2, letting FA4 leak through to the test-time fused-vs-flash comparison. On B200 runners that install flash-attn-4, this caused test_dpa_qkv_layout_thd to compare FusedAttention against an FA4 output whose padded positions contain garbage, producing 48 numerics failures in L3_pytorch_FA_versions_test--B200_1GPU. The log message already claimed FA4 would be disabled — this change makes the code match the message: set use_flash_attention_4 = False alongside use_flash_attention_2 when pad_between_seqs is True. FA3 continues to support pad_between_seqs via seqused_q/seqused_k. Signed-off-by: Sudhakar Singh <[email protected]> * [QA] Fix cutlass-dsl utils shadow in FA versions test FA4 install brings in nvidia-cutlass-dsl, whose `import cutlass` adds cutlass/base_dsl/ to sys.path. That directory contains a utils/ package that shadows tests/pytorch/utils.py, breaking collection of test_attention_with_cp.py with: ImportError: cannot import name 'ModelConfig' from 'utils' Prepend $TE_PATH/tests/pytorch to PYTHONPATH so the local utils.py is always resolved first, regardless of what FA4 dependencies install. Signed-off-by: Sudhakar Singh <[email protected]> * merge conflicts with main Signed-off-by: Sudhakar Singh <[email protected]> * fixes after merging with flash_attn_pad_bw_seqs branchj Signed-off-by: Sudhakar Singh <[email protected]> * skip tests which OOM in deterministic+backward+hopper+large_configs as its a known cudnn issue Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch][CP] Replace Python-loop THD reorder with kernel-backed permutation cache The reorder_seq_chunks_{before,after}_a2a_*_thd functions used a Python for-loop of torch.arange calls (2*cp_size*batch iterations) to build index tensors, which dominated wall-clock time at high batch counts. Replace with the existing thd_get_partitioned_indices CUDA kernel (one call per CP rank) plus a permutation cache keyed on (cu_seqlens, cp_size). This collapses thousands of tiny elementwise kernel launches into a handful of kernel calls on first use, then a dict lookup thereafter. Rename to reorder_thd_sequences_to_{rank_sharded,contiguous} since these are used by both A2A and AllGather CP paths, not just A2A. Measured speedups (cp=2, bf16, H100x2, 50 iters): cp_thd_2 (B=16): a2a 10.6x, all_gather 4.5x cp_thd_3 (B=8): a2a 5.3x, all_gather 3.3x bariamis_8k (B=2): a2a 3.0x, all_gather 2.2x bariamis_262k (B=2, S=262k): a2a 1.0x, all_gather 12.0x Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch][CP] Fix AllGather SBHD forward: set cu_seqlens_kv_per_step The THD AllGather restructuring left cu_seqlens_kv_per_step[i] as None for the SBHD/BSHD path, causing a TypeError in fused_attn_fwd which expects a Tensor. Build it from batch_size and the KV slice length, matching what the original prepare_outputs helper provided. Signed-off-by: Sudhakar Singh <[email protected]> * make cp det and nondet tests run in parallel whenever possible Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch][CP] Fix THD AllGather forward stream race on k_ag/v_ag The non-THD branch of AttnFuncWithCPAndKVAllGather.forward calls cp_stream.wait_stream(current_stream) after preparing k_ag/v_ag, but the THD branch (added by this PR) was missing that synchronization. Step 1 of the per-step attention loop launches on cp_stream and reads k_ag/v_ag, so without the wait it can race against the AllGather + reorder_thd_sequences_to_contiguous writes still in flight on the current stream. The race is masked at low cp_size because the reorder finishes before cp_stream actually starts reading. At cp=8 with the largest THD config (bucket128k: B=3, S_max=131072) the reorder takes long enough to outlast cp_stream's launch, and the resulting partially-written k_ag/v_ag feeds garbage values into cuDNN, surfacing async as cudaErrorIllegalInstruction. Move the wait_stream call out of the non-THD else-branch so both paths get the synchronization unconditionally. Verified that bucket128k all_gather @ cp=8 now runs cleanly (360.72 ms/iter). Identified by Greptile review (P1, discussion_r3170978192). Signed-off-by: Sudhakar Singh <[email protected]> * Add THD + FlashAttention v3 support to AllGather CP backend The all-gather KV buffer is always padded (dual-chunk boundaries at multiples of 2*cp_size), so per-step cu_seqlens must encode both tensor offsets and visibility limits. FA2's single cu_seqlens cannot serve both roles — a cumsum of visible counts gives wrong tensor offsets for every sequence after the first when visibility < allocation (causal/SWA). FA3's seqused_k resolves this: cu_seqlens_kv carries padded offsets, seqused_k carries visible counts. This is always needed for THD+AG regardless of the user-facing pad_between_seqs flag, which controls a different concept (caller-side padding in P2P/A2A). Changes: - Forward/backward: compute seqused_q/seqused_k and override cu_seqlens with padded values for all THD+AG+FA3 paths (not gated on pad_between_seqs) - Forward: extend KV visibility trimming and window_size adjustment to non-causal SWA (FA3 bottom_right alignment requires matching offset encoding) - Backward: use original window_size (ctx.window_size) for KV visibility instead of adjusted window_size_per_step - Add runtime assertion: THD+AG requires FA3 or FusedAttention - Test: replace unconditional THD+AG skip with FA3 availability gate - Docstring: explain offset/visibility problem, backend comparison, and why AG padding differs from the pad_between_seqs flag Signed-off-by: Sudhakar Singh <[email protected]> * Refactor AG THD window logic into shared get_kv_seq_info_after_all_gather The THD forward and backward paths duplicated the KV range and window adjustment computations that get_kv_seq_info_after_all_gather already implements for bshd/sbhd. This commit replaces the inline THD logic with calls to the shared function, reducing ~50 lines of duplication. THD differs from bshd in one key way: bshd slices the KV buffer (max_seqlen_kv = seq_end - seq_start), while THD passes the full buffer and controls visibility via cu_seqlens/seqused_k (max_seqlen_kv = seq_end). The function returns (kv_range, adjusted_window), and each caller derives max_seqlen_kv appropriately. Validated: 64 passed, 80 skipped, 0 failed (full AG suite, test_essential=False). Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch][CP] Address PR 2829 self-review: clarify THD mask/cu_seqlens PR 2829 (THD support for AllGather CP) had three self-review threads: 1. "reorder_seq_chunks_*_a2a_*_thd is no longer a2a-specific" — already addressed in commit 89b1066d (functions are now reorder_thd_sequences_to_{contiguous,rank_sharded}, used by both A2A and AllGather). 2. "Why only for non-THD do we add *_bottom_right masks?" — the code no longer guards the upgrade on qkv_format, so THD already gets bottom_right. Add a comment near the upgrade site spelling out *why* THD needs bottom_right too (Q chunk shorter than visible KV after AG; alignment is Q[0] -> KV[kv_len - q_len], identical reasoning to bshd/sbhd). This blocks future readers from re-introducing a THD-exclusion guard. 3. "Non-causal KV cu_seqlens should still respect CP padding boundaries" — the SWA-right-window addition (commit 89b1066d) commented the causal/SWA path but not the non-causal default. Add a comment at the non-causal initialisation (cu_seqlens_kv_original.clone()) explaining that visibility (cu_seqlens) and tensor layout (cu_seqlens_kv_padded) are intentionally separate: the kernel receives padded offsets via cu_seqlens_kv_padded_ in fused_attn_fwd/bwd, so cu_seqlens_kv carrying unpadded counts is correct, not a bug. Comments-only change; no behavioural impact. Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Fused thd_reorder kernel + sync-free CP THD reorder The CP THD load-balancing reorder built its gather/scatter permutation via tex.thd_get_partitioned_indices and then advanced-indexed in Python, with a per-call cu_seqlens.tolist() D2H sync to size the permutation. That sync serialized the host against the stream and left the GPU ~41% idle in the a2a THD path (nvbug 6179415), so THD ran well behind BSHD under CP even though both feed cu_seqlens to the same attention kernels. Replace it with a single fused CUDA kernel (nvte_cp_thd_reorder): one warp per token, float4-vectorized copy, dual-chunk source index computed on-device from cu_seqlens. total_tokens comes from x.shape so the launch needs no D2H sync. The kernel reuses the existing common THD helpers (binary_search and the factored thd_partition_src_index) rather than duplicating index math. Bit-identical to the Python reorder; 4-12x faster on the reorder itself and removes the GPU bubble (THD/BSHD a2a gap at cp2/nseg1 closes from ~1.5x to ~1.15x, the remainder being intrinsic cuDNN varlen-vs-dense kernel cost). _get_thd_reorder_perms is now sync-free (total_tokens passed in from x.shape; the .item() fallback is dead/defensive only). Also adds temporary, env-gated (NVTE_CP_PROFILE / NVTE_CP_NVTX) CUDA-event and NVTX instrumentation in context_parallel.py used to locate this bubble; inert unless enabled and slated for removal in PR cleanup. Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Sync-free thd_valid_copy kernel for AllGather CP THD fwd/bwd The AllGather CP THD path copied each step's valid output/dQ rows into the accumulator with a per-batch Python loop that called .item() twice per segment to read the padded start and valid length. Those D2H syncs stalled the host so it could not run ahead to keep the NCCL all-gather overlapped with compute: at cp2/nseg1/8k the dQ-copy loop alone occupied ~63% of the backward wall (mostly idle), leaving THD ~1.55x BSHD even though THD does only ~1.14x the GPU work (overlap 1.57x vs BSHD's 2.14x). Replace both loops with a single fused kernel (nvte_cp_thd_valid_copy): warp-per-token, float4-vectorized, segment found by binary_search on cu_seqlens_padded, copies row t iff its local offset is in [0, valid_len). The local>=0 guard is required because step-1 chunk starts are shifted past earlier tokens (cu_seqlens_padded[:-1] += chunk_size), so a token before a segment's padded start must be skipped rather than clobbering an already-written first-chunk row. total_tokens comes from x.shape, so no D2H. Reuses the existing binary_search helper; mirrors thd_reorder_kernel. Bit-identical to the Python loop (incl. the shifted step-1 layout). Restores comm/compute overlap (1.57x -> 1.78x): all_gather cp2/nseg1 8k THD/BSHD 1.55x -> 1.35x, and at realistic multi-sequence packing THD now beats BSHD (nseg=4 0.97x, nseg=8 0.94x). All CP all_gather THD tests pass. Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Fix FA3 all_gather THD allocator-reuse race in fused reorder The fused tex.thd_reorder (2dc5c15c) runs on the main stream and, unlike the index_select it replaced, its buffers are not allocator/stream-tracked. Its faster sync-free execution lets the host run ahead into the 2-stream per-step loop (step 1 on cp_stream), and PyTorch's caching allocator can recycle a block the reorder is still using -> cudaErrorIllegalAddress for FA3 + all_gather + THD on larger-token packs (bucket32k/64k/128k, mixed32k). Reproduces serially; masked by CUDA_LAUNCH_BLOCKING and PYTORCH_NO_CUDA_MEMORY_CACHING, confirming an allocator-reuse race. FusedAttention AG and FA3 a2a are unaffected; pre-kernel TE ran these FA3 AG configs fine, so this is a regression from the kernel work. Drain the main stream after the AG reorder, before the per-step loop allocates. cp_stream already waits on main, so the cp_stream-vs-main per-step overlap is preserved; cost is ~0.5-4.3% (one sync/forward, before the compute loop) vs the 11-30% the kernels recovered. Env-gated AG_REORDER_SYNC (default on) for A/B. A finer allocator/event fix (needs a memory-snapshot trace to pin the block, likely FA3-internal scratch) can replace this conservative drain later. Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Serialize FA3 AG calls on GPU FA3 hopper allocates per-call scheduler workspace internally, including semaphores that are recorded by the CUDA caching allocator only on the allocation stream. The THD all_gather CP path can launch consecutive FA3 per-step calls on separate streams, allowing a later step to reuse a just-freed workspace block while the prior step's scheduler kernel is still active, which shows up as illegal memory access under overlap. Replace the interim host-side reorder synchronize with FA3-scoped GPU stream waits between consecutive per-step forward and backward calls. This preserves the FusedAttention overlap path and avoids a host drain while preventing overlapping FA3 internal workspace lifetimes. Tests: - CUDA_VISIBLE_DEVICES=0,1 NVTE_BATCH_MHA_P2P_COMM=1 TE_PATH=/perfhome/llms/repos/te_repos/ag_thd_swa/TransformerEngine python3 -m pytest tests/pytorch/attention/test_attention_with_cp.py -k 'fused_attention and thd and all_gather and not fp8 and not bias' -q - CUDA_VISIBLE_DEVICES=0,1 NVTE_BATCH_MHA_P2P_COMM=1 torchrun --nproc-per-node=2 --master-port=47151 /perfhome/llms/repos/te_repos/ag_thd_swa/cp_bench/tests/pytorch/attention/run_attention_with_cp.py dtype=bf16 model=bucket32k qkv_format=thd kernel_backend=FlashAttention cp_comm_type=all_gather benchmark=0 log_level=WARNING thd_seqlen_pattern=24576,28672,30720,32768 Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Avoid D2H sync in THD max-logit mask The THD max-logit path built the valid-token mask with per-sequence .item() calls when cu_seqlens_q_padded was provided. In AllGather CP this is on the training hot path and blocks the host for every batch entry. Build the same interval mask on GPU using scatter_add over start/end deltas. This preserves support for non-zero padded offsets while avoiding the device-to-host synchronizations called out in PR review. Tests: - python3 -m black --check transformer_engine/pytorch/cpp_extensions/fused_attn.py - CUDA_VISIBLE_DEVICES=0,1 NVTE_BATCH_MHA_P2P_COMM=1 TE_PATH=/perfhome/llms/repos/te_repos/ag_thd_swa/TransformerEngine python3 -m pytest tests/pytorch/attention/test_attention_with_cp.py -k 'fused_attention and thd and all_gather and not fp8 and not bias' -q Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address THD AG review and lint issues Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Add THD helper kernel tests Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Clean up THD AG review comments Remove redundant runtime checks and stale local-experiment wording from the THD AllGather path because the PR review requested keeping this path focused on existing support gates and code behavior. Keep the non-essential CP matrix opt-in via NVTE_TEST_ESSENTIAL so offline validation does not require source edits. Signed-off-by: Sudhakar Singh <[email protected]> * [PyTorch] Address THD AG review follow-ups Resolve review comments for PR 2829 by tightening the THD all_gather output shape, renaming the new THD CP helper bindings, removing the unrelated pybind helper extraction, and aligning the FP8 t3hd aux handling with the post-FP8DS code path. Also clean up the THD CP test skip logic and remove an unnecessary dtype conversion from the THD max-logit mask construction. Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Remove duplicate AllGather padding assert The AllGather CP path already has a support assert that rejects padding masks for non-THD inputs. Keep the earlier THD-specific padding requirement, and rely on the later AllGather support assert for the non-THD padding case so the check is not duplicated. Signed-off-by: Sudhakar Singh <[email protected]> * Address THD CP review cleanup Signed-off-by: Sudhakar Singh <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clarify THD CP helper API names Rename the new THD CP reorder helpers so the source and destination layouts are encoded in the API names instead of a direction boolean. Also rename the valid-token copy helper to describe its per-split to rank-local accumulator role. Guard copy-valid tokens that precede the first padded THD offset before indexing shared cu_seqlens arrays; later split offsets can legitimately leave those token positions outside any valid sequence range. Signed-off-by: Sudhakar Singh <[email protected]> (cherry picked from commit 55f9f18747510d92791ef2650c65d93d9d90c27c) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update comment to clarify padding exclusion remove confusing info Signed-off-by: Sudhakar Singh <[email protected]> --------- Signed-off-by: Sudhakar Singh <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Avoid unpickling the extra state if not needed Signed-off-by: Przemek Tredak <[email protected]> * Remove pytorch changes from common and have better names Signed-off-by: ksivamani <[email protected]> * Fix envvar for tests Signed-off-by: ksivamani <[email protected]> * Fix tests Signed-off-by: Kirthi Shankar Sivamani <[email protected]> * Fix Signed-off-by: Kirthi Shankar Sivamani <[email protected]> --------- Signed-off-by: Przemek Tredak <[email protected]> Signed-off-by: Kirthi Shankar Sivamani <[email protected]> Signed-off-by: ksivamani <[email protected]> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
Signed-off-by: Vladimir Cherepanov <[email protected]>
…ile (#3130) * [PyTorch] torch.compile: wrap pybind11 UB methods as compile-time constants; fix SP memory leak; test suite hook-up Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * [PyTorch] Replace fp8_recipe in LinearBwdArgs with pre-resolved split-accumulator booleans LinearBwdArgs stored the entire FP8 recipe object so the backward could extract fp8_gemm_dgrad.use_split_accumulator and fp8_gemm_wgrad.use_split_accumulator at GEMM time. Recipe objects hold process-group references and are not serialisable as compile-time constants, making them incompatible with torch.compile custom-op paths. Replace fp8_recipe with two plain bool fields: - dgrad_use_split_accumulator (default _2X_ACC_DGRAD) - wgrad_use_split_accumulator (default _2X_ACC_WGRAD) These are resolved once in _linear_setup_ctx and passed into the args struct, so the backward consumes scalars instead of a live recipe object. Co-Authored-By: Claude Sonnet 4.6 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reset torch.compile state in destroy_ub to avoid stale assume_constant_result get_ub_is_fp8 bakes is_fp8_ubuf() as a compile-time constant; without a reset, destroy_ub + re-init with different FP8 settings would read stale values until recompile. Only affects in-memory caches, not disk. Signed-off-by: Pawel Gadzinski <[email protected]> * Provide explicit QuantizerRoles in torch.compile custom-recipe test ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski <[email protected]> --------- Signed-off-by: Pawel Gadzinski <[email protected]> Co-authored-by: Claude Sonnet 4.6 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
* Expert Parallelism: JAX primitives + VJPs --------- Signed-off-by: Phuong Nguyen <[email protected]>
…3151) Revert "Add wheel support for Newton-Schulz method via cuSolverMp (#3004)" This reverts commit 20e185c. Signed-off-by: ksivamani <[email protected]>
…izer (#3104) * [PyTorch] torch.compile: wrap pybind11 UB methods as compile-time constants; fix SP memory leak; test suite hook-up Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * [PyTorch] Replace fp8_recipe in LinearBwdArgs with pre-resolved split-accumulator booleans LinearBwdArgs stored the entire FP8 recipe object so the backward could extract fp8_gemm_dgrad.use_split_accumulator and fp8_gemm_wgrad.use_split_accumulator at GEMM time. Recipe objects hold process-group references and are not serialisable as compile-time constants, making them incompatible with torch.compile custom-op paths. Replace fp8_recipe with two plain bool fields: - dgrad_use_split_accumulator (default _2X_ACC_DGRAD) - wgrad_use_split_accumulator (default _2X_ACC_WGRAD) These are resolved once in _linear_setup_ctx and passed into the args struct, so the backward consumes scalars instead of a live recipe object. Co-Authored-By: Claude Sonnet 4.6 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reset torch.compile state in destroy_ub to avoid stale assume_constant_result get_ub_is_fp8 bakes is_fp8_ubuf() as a compile-time constant; without a reset, destroy_ub + re-init with different FP8 settings would read stale values until recompile. Only affects in-memory caches, not disk. Signed-off-by: Pawel Gadzinski <[email protected]> * Provide explicit QuantizerRoles in torch.compile custom-recipe test ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski <[email protected]> * Carry amax reduction group on the QuantizedTensor, not its quantizer The amax reduction process group is no longer stored persistently on a module quantizer or on a tensor's quantizer. No C++ changes. - TP sequence parallel: the group is set on the input/grad-output quantizer at point of use in the fwd/bwd impls (linear, layernorm_linear, layernorm_mlp, ops basic_linear), replacing the setup-time _customize_quantizers wiring. - FSDP2: the group is stored on Float8Tensor/NVFP4Tensor (set in fsdp_pre_all_gather) and applied to a throwaway quantizer copy during the in-place re-quant (update_quantized / _set_data). - quantize() strips the group off the output tensor's quantizer so it never persists on any tensor's quantizer (breaks flatten/pickle otherwise). Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Unwrap DebugQuantizer when setting amax reduction group set_quantizer_amax_reduction_group was a no-op on a DebugQuantizer (it lacks with_amax_reduction), so with nvinspect enabled the parent quantizer never got the SP amax reduction group, breaking fp8 current scaling column-parallel sequence-parallel numerics (debug test_numerics). Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> --------- Signed-off-by: Pawel Gadzinski <[email protected]> Co-authored-by: Claude Sonnet 4.6 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
…m-mem zero-copy (#3035) * Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy Signed-off-by: Phuong Nguyen <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]>
…NS_PER_RANK (#3150) * nccl with relax num_dispatch_tokens%64!=0 Signed-off-by: Phuong Nguyen <[email protected]> * Skip EP tests/examples on nodes without NVLink Signed-off-by: Phuong Nguyen <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]>
…#3141) * Preserve fprop operands for dequantized backward override Signed-off-by: Evgeny <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add test_grouped_linear_backward_override_high_precision_forces_save_original_input test Signed-off-by: root <[email protected]> --------- Signed-off-by: Evgeny <[email protected]> Signed-off-by: root <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root <[email protected]>
* Make quantized-tensor __repr__ fake-safe under torch.compile
Under torch.compile, TE quantized-tensor __repr__ methods are invoked on
FakeTensors during AOT autograd's structured logging. The repr bodies call
self._scale_inv.item() and/or self.dequantize() (which dispatches to the raw
C++ op tex.dequantize), both of which access a FakeTensor's data pointer and
raise:
RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor,
FunctionalTensor) ...
This was the sole cause of six fp8 failures in tests/pytorch/test_torch_compile.py.
Fix: add one shared helper, safe_quantized_repr, in tensor/_quantization_helpers.py
(a safe leaf module importing only torch) that builds a metadata-only repr
string. Each data-touching __repr__ now wraps its existing body in a try/except
and falls back to the helper when the data cannot be materialized. The eager
(non-fake) repr output is unchanged; only a fallback path is added.
Wrapped reprs: Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor
and their *Storage counterparts.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
Signed-off-by: Pawel Gadzinski <[email protected]>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Make quantized __repr__ fallback universal, drop FakeTensor-specific logic
Remove the FakeTensor-specific heuristic (_is_fake_data_access_error) and the
warning path from safe_quantized_repr. The fallback is now a plain metadata-only
repr triggered by any exception while materializing data, with each attribute
access individually guarded so __repr__ never raises.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
Signed-off-by: Pawel Gadzinski <[email protected]>
---------
Signed-off-by: Pawel Gadzinski <[email protected]>
Co-authored-by: Claude Opus 4.8 <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…` with `total_recv_tokens_per_rank` placeholder (#3154) * versioning EP C configs Signed-off-by: Phuong Nguyen <[email protected]> * Rename EP prepare token_counts to recv_tokens_per_expert Signed-off-by: Phuong Nguyen <[email protected]> * Add total_recv_tokens_per_rank placeholder to nvte_ep_prepare Signed-off-by: Phuong Nguyen <[email protected]> * Adapt PyTorch EP binding to versioned nvte_ep C config API Signed-off-by: Phuong Nguyen <[email protected]> * Rename EP group config max_num_sms to num_comm_sms Signed-off-by: Phuong Nguyen <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]>
… CUBLAS GGEMM heuristics (#3143) * support in grouped linear and relevant tests Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Unecessary details remove Removed details about FP8 current scaling methods. Signed-off-by: vthumbe1503 <[email protected]> * fix grouped linear module's grouped tensor path Signed-off-by: Varun Thumbe <[email protected]> * allow more current scaling use-cases.. block nvfp4+rht+single grouped weight being cuda graphable Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * some minor comment fixing Signed-off-by: Varun Thumbe <[email protected]> * fix heuristics Signed-off-by: Varun Thumbe <[email protected]> * only dealyed scaling skip in failure comment Signed-off-by: vthumbe1503 <[email protected]> * address review comment Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix for other 2 nvte APIs Signed-off-by: Varun Thumbe <[email protected]> * fix m and n Signed-off-by: Varun Thumbe <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…n megatron. (#2898) * add support for THD CUDA graph Signed-off-by: HaochenYuan <[email protected]> * modify comment Signed-off-by: HaochenYuan <[email protected]> * address @timmoon10: drop FAv2-bwd alloc gate, rely on THD tail zero-fill Signed-off-by: HaochenYuan <[email protected]> * Support graph-safe MoE aux loss token count Signed-off-by: HaochenYuan <[email protected]> * add graph guard for one zero fill Signed-off-by: HaochenYuan <[email protected]> * remove redundant zero-fill Signed-off-by: HaochenYuan <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename & remove prelude kernel Signed-off-by: HaochenYuan <[email protected]> * Update warp reduction function Co-authored-by: Tim Moon <[email protected]> Signed-off-by: Tim Moon <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: HaochenYuan <[email protected]> Signed-off-by: Tim Moon <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <[email protected]>
…ling (#3135)
Implements grouped-tensor quantize for the FP8 1D (1x128) and 2D (128x128)
block-scaling recipes in row-wise (RW), column-wise (CW) and BOTH quantization
directions. A single CUDA kernel launch walks 128x128 tiles across every tensor
in the group, with each CTA decoding its owning tensor from the device-side
GroupedTensor metadata with (N, R, K) shapes. Supports SAME_BOTH_DIMS (all
tensors identical) and VARYING_FIRST_DIM (constant K, varying R) shape
representations.
Three kernels share the dispatcher in group_quantize_blockwise_{1d,2d}:
- group_block_scaled_1d_rw_kernel: RW-only dispatch; 8 threads/row, reads
global memory directly into vec-16 registers; bypasses TMA since the
shared-memory roundtrip and ptx::mbarrier do not buy anything without
re-use in the CW path.
- group_block_scaled_1d_tma_kernel: CW-only and BOTH dispatch. TMA bulk-load
fills shared memory input cache. BOTH runs an RW pass (8 threads/row,
vec-16 read from shared memory) then a CW pass; CW-only skips the RW
pass. The CW pass uses 4 t/col with 32-row reg_data and two column passes
in the BOTH instantiation (keeps the per-thread register footprint under
the sm_90 3-CTAs/SM threshold) and 2 t/col with 64-row reg_data in the
CW-only instantiation (avoids doubling the smem-load bank-conflict
footprint that 4 t/col would introduce).
- group_block_scaled_2d_tma_kernel: RW-only, CW-only and BOTH dispatch. TMA
bulk-load fills shared memory input cache. Pass 1 stages 8 IVecs/thread
in registers while computing the per-tile scalar amax. Pass 2 quantizes
from registers, emits row-wise output, stages column-wise output to the
shared memory transpose staging buffer, then drains smem_T to global
memory.
Per-expert scale offsets:
- 1D RW: closed-form O(1) for both SAME_BOTH_DIMS and VARYING_FIRST_DIM
(each M_i is a multiple of kTileDim=128, hence of kScaleColAlign=4, so
DIVUP_TO_MULTIPLE collapses and the prefix sum reduces to a single
tensor_offsets_ptr[tensor_id]/K load).
- 2D CW: closed-form O(1) for SAME_BOTH_DIMS; CTA-cooperative warp-shuffle
prefix sum for VARYING_FIRST_DIM (non-linear DIVUP_TO_MULTIPLE on
blocks_y_t prevents a closed form). The cooperative reduction uses the
existing warp_allreduce_sum helper from common/utils.cuh.
Dequantize and bias-gradient (bgrad):
- group_dequantize_fp8_blockwise.cuh: kernels for all four modes
(1D/2D x rowwise/columnwise), inverting the per-expert layouts the
quantize kernels write.
- bgrad_group_quantize accepts Float8Block quantizers and computes dbias
per-tile column-partial in-kernel (mirroring MXFP8); reduced per expert
via the existing common::grouped_reduce_dbias.
Scale constraints: the fused grouped FP8BS path supports only unconstrained
FP32 scales (Float8BlockQuantizer::create_grouped_tensor rejects
force_pow_2_scales=True). Power-of-2 scales remain available on the
non-grouped/unfused split-quantize path used for Blackwell MXFP8 emulation.
Tests: existing parametrized grouped quantize / dequantize / bgrad tests
in test_grouped_tensor.py cover MXFP8, NVFP4, FP8 current scaling and the
newly-added FP8 block scaling recipe. tests/cpp/operator/
test_cast_float8blockwise_grouped.cu adds 72 C++ unit-test cases over
uniform/jagged shapes, all four (BD x direction) modes, K in {128, 256,
512}, and CUDA-graph capture coverage.
Kernels are gated to Hopper (sm_90) at the host dispatcher (cuBlasLt
grouped GEMM supports FP8 block-scaling only on Hopper).
JAX integration is intentionally left out of scope and deferred to a
follow-up PR.
Resolves #2525
Signed-off-by: Alp Dener <[email protected]>
…L2 Jax dist (#3159) * Keep the routing map format alive Signed-off-by: Kshitij Lakhani <[email protected]> * Fix incorrectly launched multi process EP tests in L2 Jax instead of L2 jax dist Signed-off-by: Kshitij Lakhani <[email protected]> --------- Signed-off-by: Kshitij Lakhani <[email protected]>
Signed-off-by: Santosh Bhavani <[email protected]>
* skip tests on hopper Signed-off-by: Varun Thumbe <[email protected]> * Update qa/L1_pytorch_mcore_fsdp_integration/test.sh Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…ests (#3174) fixing Blackwell skip condition for grouped FP8 block-scaling tests in C++ Signed-off-by: Alp Dener <[email protected]>
* disable 9.23.0/.1 for mxfp8 attention Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Charlene Yang <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…for QuantizedTensor (#3172) * fix torch function Signed-off-by: Varun Thumbe <[email protected]> * make fp bs dequantize autograd aware + plus test for autograd flow of QuantizedTensor Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Varun Thumbe <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [PyTorch] Make tensorless quantizers opaque value objects for torch.compile Give tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling, NVFP4) value-object semantics so torch.compile can treat them as baked-in constants: - Add opt-in value identity to the base Quantizer (_value_fields / _value_key / __eq__ / __hash__). Quantizers holding live tensors (delayed-scaling Float8Quantizer) and custom quantizers keep identity semantics. - New transformer_engine/pytorch/dynamo.py houses the torch.compile glue: __fx_repr__, value-key reconstruction and register_value_opaque_quantizer (gracefully a no-op on PyTorch builds without the opaque-object API). - Register the four tensorless quantizers as value opaque types. Also fix CustomRecipe state caching in TransformerEngineBaseModule: set_meta_tensor now rebuilds quantizers when the CustomRecipe instance changes (e.g. nested te.autocast regions) instead of reusing the first recipe's state, since every CustomRecipe shares the CustomRecipeState type but carries its own qfactory. Move the quantizer value-object tests into tests/pytorch/test_torch_compile.py and add that file to the L0 pytorch unittest QA suite. Signed-off-by: Pawel Gadzinski <[email protected]> * [PyTorch] Drop quantizer value registry; reconstruct via __fx_repr__ globals Follow-up to the value-opaque quantizer support: - Remove the module-level _QUANTIZER_VALUE_REGISTRY (qualname -> class) and _quantizer_from_value_key. __fx_repr__ now captures the quantizer class directly in the FX globals and reconstructs via _rebuild_quantizer(cls, items), matching how PyTorch's own value opaque types (e.g. DTensor placements) reconstruct themselves. This removes global mutable state and the qualname collision risk. - Consolidate the quantizer value-object tests in test_torch_compile.py down to two functions and exercise reconstruction through the public __fx_repr__ path instead of internal helpers. Signed-off-by: Pawel Gadzinski <[email protected]> * [PyTorch] Split dynamo.py into a dynamo/ package Replace the single dynamo.py module with a dynamo/ package so the torch.compile glue can grow with a clear responsibility split across the stacked branches. This branch owns the value-opaque quantizer layer. * dynamo/quantizer_opaque.py -- register_value_opaque_quantizer and helpers * dynamo/__init__.py -- re-exports the public API so callers keep importing from transformer_engine.pytorch.dynamo unchanged Signed-off-by: Pawel Gadzinski <[email protected]> * [PyTorch] Raise in quantizer __fx_repr__ when a process group is stored A value-opaque quantizer must not carry live distributed state. Scan the quantizer attributes in __fx_repr__ and raise TypeError if any holds a torch.distributed.ProcessGroup (e.g. a non-None deprecated amax_reduction_group), so it cannot be silently baked into a torch.compile FX graph. Clarify the related comments accordingly. Signed-off-by: Pawel Gadzinski <[email protected]> * [PyTorch] Cover NVFP4 in quantizer value-object test NVFP4Quantizer is registered as a value-opaque quantizer but was missing from the value-semantics / __fx_repr__ round-trip test. Add it to _VALUE_QUANTIZERS (skipped without CUDA, which it needs to construct). Signed-off-by: Pawel Gadzinski <[email protected]> * Reject a value quantizer that carries an amax reduction group in __eq__/__hash__ The amax reduction group is excluded from the value key, so a value quantizer that stored one would compare/hash equal to a groupless one and let torch.compile reuse a graph that skips the reduction. __eq__/__hash__ now raise (mirroring __fx_repr__, which already rejects any process-group-bearing quantizer). The group should be passed per quantize call, not stored on the quantizer. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Recognize value-opaque quantizers via a class flag Add is_value_opaque_quantizer() + the _te_compile_value_opaque flag stamped at registration, so dynamo-traced code can detect registered quantizers (and fall back to eager for unregistered ones). Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Address review: narrow opaque-type except, add fullgraph test, fix nvfp4 value key - Narrow register_opaque_type except to (RuntimeError, TypeError): the API is already imported above, so ImportError/AttributeError there only mask real errors. - Add test_quantizer_value_object_fullgraph exercising torch.compile(fullgraph=True) end-to-end to verify opaque-type registration took effect. - Restore missing NVFP4Quantizer._with_random_sign_mask assignment required by _value_fields()/_value_key(). Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Restore NVFP4 rht_matrix on value-key rebuild; assert quantize round-trip _rebuild_quantizer only restores value-key fields, so a reconstructed NVFP4Quantizer was missing the derived rht_matrix tensor (not hashable, so not in the value key) and failed at copy()/quantize time. Add a _rebuild_derived_state hook (called by _rebuild_quantizer) that NVFP4Quantizer uses to rebuild rht_matrix from _with_random_sign_mask (lru_cache -> cheap). Extend test_quantizer_value_object to also quantize with the original and the rebuilt quantizer and require bit-exact results (gated on HW support), so a field the kernel needs but the value key omits can no longer slip through. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Enforce process-group rejection in _value_key, not __fx_repr__; add test Move the ProcessGroup guard out of the (overridable) __fx_repr__ into Quantizer._value_key -- the single point every value-materialization path (__eq__/__hash__/__fx_repr__) goes through -- so a custom __fx_repr__ can no longer bypass it. Generalizes the old amax-only check to any field holding a ProcessGroup. Add a test that a value quantizer carrying a live group raises. Addresses review on NVIDIA#3152. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Strengthen fullgraph test: quantize/dequantize via a custom op, not passthrough Replace the trivial pass-through fullgraph test with one that drives each production quantizer through a minimal custom op (quantize + dequantize) under torch.compile(fullgraph=True) and compares to eager -- so the opaque-type registration is actually exercised inside the graph (a graph break would make fullgraph=True raise). Op registration sits right before the test. Also drop stale comments referencing the old __fx_repr__-side process-group guard. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Clarify comments: rht_matrix_random_sign_mask_t derivation; why the opaque flag - rht_matrix_random_sign_mask_t is a device-independent int derived from _with_random_sign_mask (the device only places a throwaway tensor); fix the misleading comment. - Explain why registration uses a class attribute, not a registry set: is_value_opaque_quantizer is traced inside the compile graph and dynamo can bake a getattr constant but cannot do 'type(q) in set' on the opaque class. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Reword opaque-flag comment: self-contained, no Linear reference Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Cover is_opaque_value_type with the import-safety guard too is_opaque_value_type(cls) sat between the import guard and the register_opaque_type guard, so on a partial/experimental opaque-object build it could raise RuntimeError/TypeError and crash TE import. Move it inside the same except so the 'registration never crashes import' promise holds for both calls. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Stamp value-opaque flag only after successful registration Move the _VALUE_OPAQUE_FLAG setattr to the end of register_value_opaque_quantizer, after register_opaque_type succeeds (or the type is already opaque). Previously the flag was set up front, so is_value_opaque_quantizer reported True even when the opaque-object API was missing or registration raised, since both paths are swallowed. Eager value semantics (__eq__/__hash__/__fx_repr__) are independent of the flag, so this only tightens the predicate to mean torch actually knows the type as opaque. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Drop verbose comments around value-opaque flag stamping Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Narrow value process-group check to amax_reduction_group _check_value_has_no_process_group ran on every guard eval (via __eq__/__hash__) and scanned all of vars(self) recursively. The only attribute that can hold a ProcessGroup is the deprecated amax_reduction_group, so check it directly (O(1)) and drop the _contains_process_group helper. Same guarantee, off the hot path. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Shorten amax_reduction_group check comment Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Drop trivial value-equality boilerplate from quantizer test Remove the a==b / hash / dict-key block that just exercised Python's own dict semantics; equality and hashing are still covered by the __fx_repr__ round-trip (rebuilt == a, hash match) and the bit-exact kernel check. other_kwargs is now unused, so drop it from the parametrization and both test signatures. Co-Authored-By: Claude Opus 4.8 <[email protected]> Signed-off-by: Pawel Gadzinski <[email protected]> * Address review comments: qualname registry, import and comment cleanups - Replace the class-attribute value-opaque flag with a module-level set of class qualnames: a set of class objects is untraceable under fullgraph=True (opaque classes have no equality rule in Dynamo), but the qualname constant-folds to a plain string; also avoids falsely reporting unregistered subclasses. - Register MXFP8Quantizer right after the class like the other quantizers. - Clarify the amax_reduction_group exclusion comment in Float8CurrentScalingQuantizer._value_fields. - Restore import order in test_torch_compile.py, import NVFP4Quantizer from transformer_engine.pytorch, drop unused Float8Quantizer import. Signed-off-by: Pawel Gadzinski <[email protected]> * Derive quantizer value fields from class annotations Make the class annotations the single source of truth for what defines a quantizer's value, instead of hand-written per-class _value_fields lists: - Quantizer._value_fields is now derived from the annotations across the MRO (subsuming _BASE_VALUE_FIELDS); register_value_opaque_quantizer is the explicit opt-in and validates at import time that no annotated field is a tensor or process group. - Drop the four per-class _value_fields overrides. - Remove the deprecated amax_reduction_group annotation from Float8CurrentScalingQuantizer and NVFP4Quantizer (the attribute is still set for backward compatibility). - NVFP4: rename _with_random_sign_mask to with_random_sign_mask (annotated, matching the constructor argument), stop storing the derived rht_matrix_random_sign_mask_t in the value key and rebuild it together with rht_matrix in _rebuild_derived_state (lru-cached getters), which __init__ now also uses. copy() now propagates with_random_sign_mask. Signed-off-by: Pawel Gadzinski <[email protected]> * Drop redundant annotation-exclusion comments Signed-off-by: Pawel Gadzinski <[email protected]> * Shorten _is_value_quantizer comment Signed-off-by: Pawel Gadzinski <[email protected]> * Fold annotation walk into _value_fields, trim comments Signed-off-by: Pawel Gadzinski <[email protected]> * Simplify qualname-registry comment Signed-off-by: Pawel Gadzinski <[email protected]> * Reword qualname-registry comment for outside readers Signed-off-by: Pawel Gadzinski <[email protected]> * Fix review findings: pickle compat, subclass opt-in, cached value fields - NVFP4: stop storing with_random_sign_mask; the annotated value field is the derived (deterministic, device-independent) rht_matrix_random_sign_mask_t and rht_matrix is rebuilt from it in _rebuild_derived_state. Quantizers pickled before this PR already carry the mask in __dict__, so old checkpoints keep working (a boolean back-fill default could lie for quantizers created with with_random_sign_mask=False). - Value semantics no longer leak into unregistered subclasses: register_value_opaque_quantizer stores the field tuple on the class and _value_fields looks it up in the class's own __dict__, so a subclass must register explicitly (it previously inherited value eq/hash that ignored its unannotated fields and skipped the annotation check). - The value-field tuple is computed once at registration instead of an MRO walk per __eq__/__hash__ call (these run per compiled-function invocation via the EQUALS_MATCH guard); registration validation and field derivation now share one annotation walk (Quantizer._annotated_fields). Drop the unreachable other._value_fields() branch and do the cheap type check first in __eq__. Signed-off-by: Pawel Gadzinski <[email protected]> * Validate value-field annotations by resolved type, not annotation text Replace the substring blocklist with get_type_hints + an allowlist of value types (int/bool/float/str/enum): aliased tensor types no longer slip through and benign types whose name merely contains "Tensor" are no longer rejected. Runs once per class at import time, not in any hot path. Signed-off-by: Pawel Gadzinski <[email protected]> * Drop the amax_reduction_group fixup from the generic rebuilder _rebuild_quantizer no longer back-fills the deprecated amax_reduction_group (and loses the field_names set that existed only for that check): a rebuilt quantizer deliberately lacks the attribute, so anything that genuinely needs it fails loudly instead of silently getting None. The only unconditional readers were the two copy() methods, which now tolerate the absent field; _canonicalized_amax_reduction_group (used by the kernel only when with_amax_reduction is set) still raises AttributeError on a rebuilt quantizer, which is the intended behavior. Signed-off-by: Pawel Gadzinski <[email protected]> * Enable post-RHT amax in the NVFP4 value-object test factory The quantize kernel rejects with_rht=True without with_post_rht_amax=True (pre-RHT amax unsupported); mirror the recipe, which always sets both together. Unnoticed locally because the NVFP4 round-trip is skipped on non-NVFP4 hardware. Signed-off-by: Pawel Gadzinski <[email protected]> --------- Signed-off-by: Pawel Gadzinski <[email protected]> Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
* Remove cuDNN frontend git submodule Use nvidia-cudnn-frontend for the C++ headers and Python bindings. Keep the cuDNN library discovery helper in tree and update common, PyTorch, JAX, packaging, and test build paths. Signed-off-by: Vladimir Cherepanov <[email protected]> * Fix JAX isolated build requirements Signed-off-by: Vladimir Cherepanov <[email protected]> --------- Signed-off-by: Vladimir Cherepanov <[email protected]>
…behind cuBLAS 13.5+ (#3181) * Guard FP8 per-tensor scaling grouped GEMM on Hopper behind cuBLAS 13.5+ cuBLAS 13.4 has no SM90 grouped GEMM algorithms for the per-tensor (PER_BATCH_SCALAR_32F) FP8 scale mode, so the heuristic query fails with CUBLAS_STATUS_NOT_SUPPORTED and users get a cryptic "Unable to find suitable cuBLAS grouped GEMM algorithm" error (see #3176). The support is available starting with cuBLAS 13.5. - Add CUBLAS_FP8_TENSOR_SCALING_GROUPED_GEMM_HOPPER_VERSION (130500) and a runtime check with a clear error message on the FP8 tensor-scaling grouped GEMM path. - Enable the FP8Current grouped GEMM C++ tests on Hopper with cuBLAS 13.5+ (previously skipped as Blackwell-only, so this path had no C++ coverage on SM90). - Skip FP8 current scaling GroupedLinear tests on Hopper when cuBLAS < 13.5 instead of failing in algorithm selection. Signed-off-by: Pawel Gadzinski <[email protected]> * Fall back to legacy GroupedLinear path when fused grouped GEMM is unavailable Route around the fused GroupedTensor path instead of erroring out when the cuBLAS version is too old: grouped GEMM needs cuBLAS 13.3+ (13.4+ on Hopper), and FP8 per-tensor current scaling on Hopper needs 13.5+. Both the module (_is_grouped_tensor_path_supported) and the ops (_is_graph_safe_path_supported) checks now fall back to the legacy multi-stream path in these cases. Signed-off-by: Pawel Gadzinski <[email protected]> --------- Signed-off-by: Pawel Gadzinski <[email protected]>
* [NCCL EP] Bump nccl submodule to 2.31.0a14 (b87848fbc5) Signed-off-by: Phuong Nguyen <[email protected]> * [PyTorch][JAX] Support int32 topk_idx in NCCL EP dispatch and prepare Signed-off-by: Phuong Nguyen <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]>
… kernel (#3190) * [PyTorch] NVFP4: emit GEMM-swizzled scales in the non-RHT 2D quantize kernel Fixes a CUDA graph hang/corruption with NVFP4 weight caching. Scope: the non-RHT 2D quantize path on 128-aligned shapes, which is the default cached-weight case. Other paths (1D weight scaling, non-128-aligned, etc) keep the compact + post-quantize swizzle fallback and are left for follow-ups. - swizzle.cuh: NVFP4 128x4 GEMM-swizzled scale-index helper (byte-compatible with nvte_swizzle_scaling_factors). - quantize_transpose_nvfp4_2D_kernel: add WITH_GEMM_SWIZZLED_SCALES; write rowwise and columnwise scales at swizzled offsets; relax the compact-only guard to allow the 2D swizzled path on 128-aligned dims. - dispatch: thread with_gemm_swizzled_scales through the blockwise fallback (which already implements kSwizzledScale) instead of hardcoding false. - quantizer: is_eligible_for_2d_swizzle_fusion + shared gate so the non-RHT 2D path enables swizzled SF on 128-aligned shapes; the post-quantize swizzle fallback auto-skips when the flag is set. - tests: byte-equal SF vs nvte_swizzle_scaling_factors (rowwise/columnwise/ both, multiple shapes) + shape gate; cached-weight scale-pointer stability and CUDA graph capture/replay regression. Signed-off-by: Cael Ling <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Cael Ling <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 <[email protected]>
* Remove fused attention workspace optimization control Signed-off-by: Przemek Tredak <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * More removal Signed-off-by: Przemek Tredak <[email protected]> --------- Signed-off-by: Przemek Tredak <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Enable NVFP4 RHT amax for grouped SReLU MLP Signed-off-by: Siddhartha Raman <[email protected]> * Update tests/pytorch/test_fusible_ops.py Co-authored-by: vthumbe1503 <[email protected]> Signed-off-by: Siddhartha Raman Sundara Raman <[email protected]> * Fix indentation for tols assignment in test Signed-off-by: vthumbe1503 <[email protected]> * Guard SReLU hadamard kernel by cuDNN frontend version Signed-off-by: Siddhartha Raman Sundara Raman <[email protected]> * Update tolerance handling for quantization tests Set default tolerance values for quantization checks. Signed-off-by: vthumbe1503 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert back the tolerance diff Simplified tolerance settings for quantization tests. Signed-off-by: vthumbe1503 <[email protected]> --------- Signed-off-by: Siddhartha Raman <[email protected]> Signed-off-by: Siddhartha Raman Sundara Raman <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Co-authored-by: vthumbe1503 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [JAX] Resync onto upstream PR #3036, restore TE-EP-only MoE block
Reset 33 local commits onto phuong/ep-3-jax @ c34771d4 (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 <[email protected]>
* 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 <[email protected]>
* 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 <[email protected]>
* 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 <[email protected]>
* jax/flax: realign _MoEBlock with post-resync moe() signature
After the upstream PR #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 <[email protected]>
* 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 <[email protected]>
* jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR #3116)
Address jberchtold-nvidia's PR #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 <[email protected]>
* jax/moe: address PR #3116 review feedback (hardcode align + expand inline justifications)
Responds to jberchtold-nvidia's PR #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 #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 <[email protected]>
* 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 #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 <[email protected]>
* 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 <[email protected]>
* remove useless comments
Signed-off-by: Teddy Do <[email protected]>
* 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 #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 <[email protected]>
* jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate
Per reviewer feedback (Jaberchtold on PR #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 <[email protected]>
* jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract
Mirrors the sparse_probs NaN-sanitizer removal in fe446974: 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 <[email protected]>
* 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 <[email protected]>
* 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 <[email protected]>
* jax/moe: address TE EP alignment review feedback
Signed-off-by: Teddy Do <[email protected]>
* jax/moe: fix early topk weighting padded-slot masking
Signed-off-by: Teddy Do <[email protected]>
* jax/moe: remove unused EP mesh size
Signed-off-by: Teddy Do <[email protected]>
* jax/moe: tighten TE EP recv capacity bound
Signed-off-by: Teddy Do <[email protected]>
* jax/moe: simplify late TE EP weighting
Signed-off-by: Teddy Do <[email protected]>
* [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 <[email protected]>
---------
Signed-off-by: Teddy Do <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…t bindings (#3203)
* [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 <[email protected]>
* Add fully qualified name to softmax kernels
Signed-off-by: CarlosGomes98 <[email protected]>
* Add static fallback option, fix softmax acc_t dtype
Signed-off-by: CarlosGomes98 <[email protected]>
* add missing license
Signed-off-by: CarlosGomes98 <[email protected]>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* greptile changes
Signed-off-by: CarlosGomes98 <[email protected]>
* fix formatting, .clang-format
Signed-off-by: CarlosGomes98 <[email protected]>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* import cleanup
Signed-off-by: CarlosGomes98 <[email protected]>
* Test more columns for softmax, mr changes
Signed-off-by: CarlosGomes98 <[email protected]>
* Fix tests
Signed-off-by: Carlos Gomes <[email protected]>
---------
Signed-off-by: CarlosGomes98 <[email protected]>
Signed-off-by: Carlos Gomes <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Tim Moon <[email protected]>
* support scaled swiglu, scaled srelu and scaled clamp swiglu Signed-off-by: zhongboz <[email protected]> * vectorized loading improvement Signed-off-by: Zhongbo Zhu <[email protected]> * fix bug for backward kernel Signed-off-by: Zhongbo Zhu <[email protected]> * optimize Signed-off-by: Zhongbo Zhu <[email protected]> * fix unit test failure Signed-off-by: Zhongbo Zhu <[email protected]> * [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 <[email protected]> * resolve comments Signed-off-by: Zhongbo Zhu <[email protected]> * refactor, resolve comments Signed-off-by: Zhongbo Zhu <[email protected]> * address review comment Signed-off-by: Varun Thumbe <[email protected]> * adaptive cta to fix slow block reduce for scale grads Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * [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 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: zhongboz <[email protected]> Signed-off-by: Zhongbo Zhu <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Signed-off-by: Varun Thumbe <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 <[email protected]>
…DNN SDPA fprop (#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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
* [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 <[email protected]>
Signed-off-by: Emil Gilliam <[email protected]>
---------
Signed-off-by: Emil Gilliam <[email protected]>
Co-authored-by: Claude Fable 5 <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sudhakar Singh <[email protected]>
…… (#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 <[email protected]>
Signed-off-by: Tim Moon <[email protected]>
* fix grouped linear hang Signed-off-by: Varun Thumbe <[email protected]> * make the same change in grouped mlp as well Signed-off-by: Varun Thumbe <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]>
…GroupedLinear and fused grouped MLP (#3161) * Add optional caller-provided output/grad-input buffers to GroupedLinear module and fusible ops Signed-off-by: Phuong Nguyen <[email protected]> * [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 <[email protected]> * Write fused grouped MLP MXFP8 output and dgrad directly into caller buffers, eliminating the D2D copy + cleanup Signed-off-by: Phuong Nguyen <[email protected]> * [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 <[email protected]> * use basic_ops to track op kwargs Signed-off-by: YangFei1990 <[email protected]> * add doc and resolve comments Signed-off-by: YangFei1990 <[email protected]> * move out/dgrad_out out from the non_tensor_args Signed-off-by: YangFei1990 <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]> Signed-off-by: YangFei1990 <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: YangFei1990 <[email protected]> Co-authored-by: Fei Wu <[email protected]>
* Fix FusedAdam empty tensor handling Signed-off-by: Jingyue Wu <[email protected]> * Move empty tensor filtering into MultiTensorApply Signed-off-by: Jingyue Wu <[email protected]> --------- Signed-off-by: Jingyue Wu <[email protected]> Co-authored-by: vthumbe1503 <[email protected]>
matthiasdiener
marked this pull request as ready for review
August 27, 2026 22:01
matthiasdiener
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
August 27, 2026 22:01
ipanfilo
approved these changes
Aug 27, 2026
Co-authored-by: Aristotle Martin <[email protected]>
matthiasdiener
force-pushed
the
IFU-dev-20260803-v2.18-merge
branch
from
August 27, 2026 23:02
297351a to
0b1399e
Compare
Contributor
Author
The merges with dev were not identical, should be resolved now. |
ipanfilo
approved these changes
Aug 27, 2026
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
20 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
Squashed all ROCm-specific commits from #702 into one commit.
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: