Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
70 commits
Select commit Hold shift + click to select a range
f6ed27c
[PyTorch][torch.compile] Support for DotProductAttention
pggPL Jul 27, 2026
3053d76
Delegate the no-FP8 case of DPA's init_fp8_metadata to the base class
pggPL Jul 28, 2026
d6b285e
Run DPA eagerly for undeclared packed q/k/v instead of guessing the l…
pggPL Jul 28, 2026
1cd6704
Log backend selection in eager only; shorten comments
pggPL Jul 28, 2026
42a9e6c
Fix lint: unused import and keyword-arg-before-vararg
pggPL Jul 29, 2026
af11f8f
Compare against eager under CUDA graphs, and tighten the tolerances
pggPL Jul 29, 2026
0b64c96
Support one cu_seqlens tensor for both q and kv under torch.compile
pggPL Jul 29, 2026
7f480c6
Apply the shared-cu_seqlens workaround to the FlashAttention v4 path too
pggPL Jul 29, 2026
adc96c4
Run FP8 attention eagerly under torch.compile
pggPL Jul 29, 2026
737904f
Drive the compile tests off ModelConfig instead of a hand-picked inte…
pggPL Jul 29, 2026
9026a41
Scale the unfused comparison to the tensor, and take tolerances from …
pggPL Jul 29, 2026
bd2c506
Merge the per-backend compile tests, and give padding masks real padding
pggPL Jul 29, 2026
725931e
Run eagerly when max_seqlen has to be derived from cu_seqlens
pggPL Jul 29, 2026
b2ab7eb
Share the run-and-compare code between the compile tests
pggPL Jul 29, 2026
50d1984
Use the shared run-and-compare code in the remaining compile tests
pggPL Jul 29, 2026
f19c52d
Read the predicate's arguments by name, not by hardcoded position
pggPL Jul 29, 2026
fbc7146
Move the eager-fallback decorator to jit.py
pggPL Jul 29, 2026
6b24630
Skip FusedAttention rather than special-casing it in the compile tests
pggPL Jul 29, 2026
9a204b8
Cover the compiled path around FusedAttention
pggPL Jul 30, 2026
2188f18
Log backend selection through the no-op logger, as get_attention_back…
pggPL Jul 30, 2026
10d141b
Return the fused sub-backend as an int, in eager as well
pggPL Jul 30, 2026
93a201d
Rename eager_under_compile_if to fallback_to_eager_when
pggPL Jul 30, 2026
0866319
Fold the eager fallback into no_torch_dynamo as a `when` predicate
pggPL Jul 30, 2026
a80a41a
Cover the declared packed layouts beyond bs3hd
pggPL Jul 30, 2026
9eb1449
Let ONNX export keep its own path in get_qkv_layout
pggPL Jul 30, 2026
5c5b451
Guard the assumption the argument binding rests on
pggPL Jul 30, 2026
95030f9
Check that CUDA graphs were actually captured, and generalize a docst…
pggPL Jul 30, 2026
b0e56d4
Skip lazy compilation while already tracing
pggPL Jul 30, 2026
6f7c818
Merge upstream/main into dpa_torch_compile
pggPL Jul 30, 2026
2cf695b
Run context parallel attention eagerly
pggPL Jul 30, 2026
b89446f
Keep DotProductAttention compilable under a CUDA RNG states tracker
pggPL Aug 4, 2026
fb50154
Run checkpointed attention eagerly
pggPL Aug 4, 2026
4b1ebb2
Support KV caching under torch.compile
pggPL Aug 4, 2026
eb5ad55
Compile the fused sbh3d QKV split
pggPL Aug 4, 2026
7da1a66
Run FlashAttention 4 eagerly
pggPL Aug 4, 2026
d1beb9b
Pin the sequence lengths backend selection bakes in
pggPL Aug 4, 2026
4af8c30
Capture CUDA graphs with a KV cache
pggPL Aug 4, 2026
bd94fff
Test generation against a KV cache
pggPL Aug 4, 2026
688e10c
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 4, 2026
43eba11
Drop the imports the custom ops made unused
pggPL Aug 4, 2026
4ff5fdf
[common] Fix UE8M0 code 0 (2^-127) and code 255 (NaN) expansion in pt…
xiuhu17 Aug 5, 2026
a7aec21
[Pytorch] Add function to check if EP is bootstrapped (#3321)
YangFei1990 Aug 5, 2026
66b6d25
[JAX] Migrate EP collective-stream annotation to region-based compute…
phu0ngng Aug 6, 2026
46e0edf
[PyTorch]Fused ScaledActivation Kernel in TE Ops and reducing CPU ove…
vthumbe1503 Aug 6, 2026
dd96595
Add support for fused Q Up-Proj GEMM/RoPE/Quant. (#3303)
chaseblock Aug 7, 2026
da28058
[Pytorch] Add support for row-wise quanted input for grouped gemm (#3…
YangFei1990 Aug 7, 2026
8260f49
Use a temp directory for CMake build (#3254)
fheinecke Aug 7, 2026
c4a38bc
Rename _maybe_unshare_cu_seqlens to _unalias_cu_seqlens
pggPL Aug 7, 2026
1b495ce
Split the get_full_cu_seqlens cache bypass per reason
pggPL Aug 7, 2026
753cfeb
Merge remote-tracking branch 'upstream/main' into dpa_merge
pggPL Aug 7, 2026
731f242
[PyTorch] Preserve FP8 recompute state for inner autocast (#3284)
nvegesna-netizen Aug 10, 2026
32f7096
[JAX][Fix] Gate collective-stream compute_on on import and silence py…
phu0ngng Aug 10, 2026
598b9ea
[Bug] use_single_group_runtime_offsets arg is not supported for Scale…
vthumbe1503 Aug 10, 2026
f07a860
[Common] Fix pointer arithmatic to generate correct LDS/STS instructi…
kainzhong Aug 10, 2026
07e281f
[Common] MXFP8: Add swizzled-scale fast path for cast-only quantizati…
WanZzzzzz Aug 11, 2026
b54bb1e
Enable runtime resolution of CUDA header path for NVRTC (#3252)
fheinecke Aug 11, 2026
56a05e7
[Common][PyTorch] EP dispatch with unfused MXFP8 quantization (#3270)
phu0ngng Aug 12, 2026
734ce31
[PyTorch] Advance FusedAdam step counter for empty param groups (#3318)
adityasingh2400 Aug 12, 2026
fc1fc0e
[Common] Fix NVFP4 stochastic rounding on architectures without cvt.r…
davidkny22 Aug 12, 2026
1ad30c6
[Pytorch] Fix swizzle scale output shape in variable shape case (#3344)
YangFei1990 Aug 12, 2026
779f17d
[Common] Upgrade nccl-extensions commit (#3349)
phu0ngng Aug 12, 2026
2d80391
Improve device-init grouped linear module with single grouped weight …
zhongbozhu Aug 12, 2026
67c26e2
fix: allgather striped window_size (5,0) id says (8,0) (#3374)
andrewwhitecdw Aug 13, 2026
910be8b
fix: remove redundant second export of XLA_FLAGS in L0 (#3377)
andrewwhitecdw Aug 13, 2026
900ac82
fix: remove unused `import time` (#3364)
andrewwhitecdw Aug 13, 2026
7d0365b
Stabilize FlashAttention CI environments (#3341)
sudhakarsingh27 Aug 13, 2026
8f929e6
Merge branch 'main' of NVIDIA/TransformerEngine into dpa_torch_compile
pggPL Aug 13, 2026
76ebc3c
Skip get_fp8_recipe only while tracing in init_fp8_metadata
pggPL Aug 13, 2026
075bce8
Skip the max_seqlen bake-in under ONNX export
pggPL Aug 17, 2026
b3e9949
Make undeclared non-contiguous q/k/v contiguous under compile
pggPL Aug 17, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rdparty/nccl-extensions
Submodule nccl-extensions updated 109 files
Original file line number Diff line number Diff line change
Expand Up @@ -13,21 +13,21 @@

Example:

python benchmarks/linear/benchmark_graph_safe_grouped_linear.py
python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py

Forward-only:

python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --fwd-only
python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --fwd-only

Nsight Systems:

(optionally: unset DEBUGINFOD_URLS)

nsys profile \
--output=./benchmarks/linear/graph_safe_grouped_linear_mxfp8 \
--output=./benchmarks/linear/graph_safe_grouped_mlp_mxfp8 \
--force-overwrite true \
--trace=cuda,nvtx,cudnn,cublas \
python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --profile
python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --profile
"""

# Match the Qwen MXFP8 SFT launch toggles before importing TE.
Expand Down
Loading
Loading