Skip to content

Gpt oss jax Changes for me to look at - #1

Open
vthumbe1503 wants to merge 63 commits into
users/vthumbe/gpt_oss_swiglu_integrationfrom
gpt-oss-jax
Open

Gpt oss jax Changes for me to look at#1
vthumbe1503 wants to merge 63 commits into
users/vthumbe/gpt_oss_swiglu_integrationfrom
gpt-oss-jax

Conversation

@vthumbe1503

@vthumbe1503 vthumbe1503 commented Sep 22, 2025

Copy link
Copy Markdown
Owner

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

phu0ngng and others added 5 commits September 22, 2025 12:58
@phu0ngng
[JAX] Remove import jax.extend.ffi (NVIDIA#2193)
57b4d7b 
* remove import jax.extend.ffi

Signed-off-by: Phuong Nguyen <[email protected]>

---------

Signed-off-by: Phuong Nguyen <[email protected]>
@cyanguwa @cuichenx @pre-commit-ci
[PyTorch] Add sink attention support from cuDNN (NVIDIA#2148)
5e4e0b2 
* first draft; debug plan failure

Signed-off-by: Charlene Yang <[email protected]>

* debug uid error

Signed-off-by: Charlene Yang <[email protected]>

* tweak params

Signed-off-by: Charlene Yang <[email protected]>

* add grad in output

Signed-off-by: Charlene Yang <[email protected]>

* clean up prints

Signed-off-by: Charlene Yang <[email protected]>

* fix prints in test

Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* address review comments

Signed-off-by: Charlene Yang <[email protected]>

* fix unfused grad; add softmax_type; add sink to bwd

Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* fix padding mask; add swa tests; remove requires_grad for off-by-one

Signed-off-by: Charlene Yang <[email protected]>

* update FE

Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* Apply 1 suggestion(s) to 1 file(s)

Co-authored-by: Chen Cui <[email protected]>
Signed-off-by: Charlene Yang <[email protected]>

* fix indent

Signed-off-by: Charlene Yang <[email protected]>

* fix non-determinism and shapes

Signed-off-by: Charlene Yang <[email protected]>

* clean up prints

Signed-off-by: Charlene Yang <[email protected]>

* add GQA

Signed-off-by: Charlene Yang <[email protected]>

* add CP A2A; dq/dk mismatches

Signed-off-by: Charlene Yang <[email protected]>

* fix CP A2A; need cleaner solution

Signed-off-by: Charlene Yang <[email protected]>

* fix CP A2A; pending cudnn kernel change

Signed-off-by: Charlene Yang <[email protected]>

* minor fixes

Signed-off-by: Charlene Yang <[email protected]>

* fix world size in unit test; avoid thd format

Signed-off-by: Charlene Yang <[email protected]>

* fix kernel_backend, dtype in unit test; fix head_dim for FP8 Hopper

Signed-off-by: Charlene Yang <[email protected]>

* fix thd logic

Signed-off-by: Charlene Yang <[email protected]>

* fix fp8 context

Signed-off-by: Charlene Yang <[email protected]>

* tweak CP logging

Signed-off-by: Charlene Yang <[email protected]>

* allow no_mask/padding for SWA(left,0)

Signed-off-by: Charlene Yang <[email protected]>

* Revert "allow no_mask/padding for SWA(left,0)"

This reverts commit 08b4ccc67a08b6882080b06aa715f541bb832aca.

Signed-off-by: Charlene Yang <[email protected]>

* add softmax_type to Jax

Signed-off-by: Charlene Yang <[email protected]>

* add cuDNN version control

Signed-off-by: Charlene Yang <[email protected]>

* prettify tests

Signed-off-by: Charlene Yang <[email protected]>

* skip 9.13 for MLA, non 192/128

Signed-off-by: Charlene Yang <[email protected]>

* rename compare_with_error

Signed-off-by: Charlene Yang <[email protected]>

* small cleanups and improvements

Signed-off-by: Charlene Yang <[email protected]>

* fix minor CI failures

Signed-off-by: Charlene Yang <[email protected]>

* force sink/dsink to be float32

Signed-off-by: Charlene Yang <[email protected]>

* switch FE to GH FE

Signed-off-by: Charlene Yang <[email protected]>

* return to GH TE main FE commit

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

* update FE to 1.14.1

Signed-off-by: Charlene Yang <[email protected]>

* clean up before CI

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

* fix lint

Signed-off-by: Charlene Yang <[email protected]>

* bump up cudnn version

Signed-off-by: Charlene Yang <[email protected]>

* add backend selection guard for unit tests

Signed-off-by: Charlene Yang <[email protected]>

* add docstring for softmax type enums in C

Signed-off-by: Charlene Yang <[email protected]>

---------

Signed-off-by: Charlene Yang <[email protected]>
Co-authored-by: Chen Cui <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@shengfangd
[QA] Add pytest xml report for all tests in qa folder that use pytest (…
2db20a6 
…NVIDIA#2169)

* Add pytest xml report for debug unittest and onnx unittest, and remove the duplicated test line in qa/L0_pytorch_debug_unittest/test.sh

---------

Signed-off-by: erindai <[email protected]>
@mingxu1067 @phu0ngng
[JAX] Local-Amax for Current-Scaling (NVIDIA#2183)
a92a0ad 
* Adding Amax Primitive and related args.

Signed-off-by: Ming Huang <[email protected]>

* Enable local-amax for current-scaling and optionally run AR aross FSDP/TP/SP.

Signed-off-by: Ming Huang <[email protected]>

* Adding doc for Amax Primitive.

Signed-off-by: Ming Huang <[email protected]>

* Fix the function name conflict.

Signed-off-by: Ming Huang <[email protected]>

* Modification as feedback suggested.

Signed-off-by: Ming Huang <[email protected]>

* Fix errors from lint.

Signed-off-by: Ming Huang <[email protected]>

* Fix the wrong amax-scope in the bwd.

Signed-off-by: Ming Huang <[email protected]>

* Added more description for amax-scope

Signed-off-by: Ming Huang <[email protected]>

* Fix the wrong attribute name.

Signed-off-by: Ming Huang <[email protected]>

* Keep dim for AmaxCalcuation.

Signed-off-by: Ming Huang <[email protected]>

* Remove keepDim and add shardy_rule

Signed-off-by: Ming Huang <[email protected]>

* Fix shardy_rule

Signed-off-by: Ming Huang <[email protected]>

* Remove extra-collective bytes from ref_coll_count due to local amax.

Signed-off-by: Ming Huang <[email protected]>

---------

Signed-off-by: Ming Huang <[email protected]>
Signed-off-by: Ming-Xu Huang <[email protected]>
Co-authored-by: Phuong Nguyen <[email protected]>
@phu0ngng
[JAX] Restore Shardy Rule with CompoundFactor (NVIDIA#2167)
3f875fb 
* Rework shardy rules

* WAR for compound factor=1

Signed-off-by: Phuong Nguyen <[email protected]>

---------

Signed-off-by: Phuong Nguyen <[email protected]>
@vthumbe1503
vthumbe1503 force-pushed the gpt-oss-jax branch 6 times, most recently from e81a0e1 to 6675779 Compare September 24, 2025 00:47
@vthumbe1503
JAX integration changes
0c17c7e 
Signed-off-by: Varun Thumbe <[email protected]>
vthumbe1503 and others added 17 commits September 23, 2025 18:41
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
90e070c 
Signed-off-by: vthumbe1503 <[email protected]>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
66c7086 
for more information, see https://pre-commit.ci
@vthumbe1503
revert line break
af19dbf 
Signed-off-by: vthumbe1503 <[email protected]>
@vthumbe1503
revert line break
4f29915 
Add documentation for quantization function parameters and return value.

Signed-off-by: vthumbe1503 <[email protected]>
@vthumbe1503
missed adding oss swiglu to nvte enum in common
24828f3 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
fix jax linting errors
19410b6 
Signed-off-by: Varun Thumbe <[email protected]>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
5480d29 
for more information, see https://pre-commit.ci
@vthumbe1503
fix jax linting errors
7a917ea 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
revert multi_gpu_encoder change
53dd179 
Signed-off-by: Varun Thumbe <[email protected]>
@phu0ngng
[JAX] Update JAX version requirement in pyproject.toml (NVIDIA#2197)
afd15a1 
update jax requirements

Signed-off-by: Phuong Nguyen <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
d048807
@vthumbe1503
fix flax integration bug
3bfae54 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'gpt-oss-jax' of github.com:vthumbe1503/TransformerEngin…
9c60c47 
…e into gpt-oss-jax

Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
fix linting error
38382dc 
Signed-off-by: Varun Thumbe <[email protected]>
@pggPL
[PyTorch] Unpin version of onnxscript and onnxruntime (NVIDIA#2202)
9e72796 
* fix

Signed-off-by: Pawel Gadzinski <[email protected]>

* fix

Signed-off-by: Pawel Gadzinski <[email protected]>

* fix

Signed-off-by: Pawel Gadzinski <[email protected]>

* fix

Signed-off-by: Pawel Gadzinski <[email protected]>

---------

Signed-off-by: Pawel Gadzinski <[email protected]>
@vthumbe1503
bug fixed in other branch and not here
c7ef078 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
c39ab8d
jberchtold-nvidia and others added 29 commits September 30, 2025 16:53
@jberchtold-nvidia @nvjax
[JAX] Load modules during initialize for Norm and Act primitives (NVI…
ce18bee 
…DIA#2219)

Load modules during initialize

Signed-off-by: Jeremy Berchtold <[email protected]>
Co-authored-by: JAX Toolbox <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
b7df6b6 
Signed-off-by: vthumbe1503 <[email protected]>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
4f41c1b 
for more information, see https://pre-commit.ci
@negvet @pre-commit-ci @timmoon10
[PyTorch] Quantizer as API (NVIDIA#2039)
7022d50 
* Introduce QuantizerBase

Signed-off-by: Evgeny <[email protected]>

* Expose as a first-class API

Signed-off-by: Evgeny <[email protected]>

* Undo QuantizerBase

Signed-off-by: Evgeny <[email protected]>

* Make Quantizer a base class without implementations

Signed-off-by: Evgeny <[email protected]>

* Support CustomRecipe and CustomRecipeState

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Resolving comments: quantize impl, num_quantizers, defaults

Signed-off-by: Evgeny <[email protected]>

* Quantizer factories

Signed-off-by: Evgeny <[email protected]>

* Add tests

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* QuantizedTensorBase _get_quantizer() + quantize_()

Signed-off-by: Evgeny <[email protected]>

* Experimental note + LayerNormMLP fix

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tensor._internal -> tensor.base

Signed-off-by: Evgeny <[email protected]>

* Expose

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Minor import fix

Signed-off-by: Evgeny <[email protected]>

* Single quantizer factory with roles

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* More context for qfactory, fwd/bwd_roles

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Minor

Signed-off-by: Evgeny <[email protected]>

* Rename *Base -> *Storage quantized tensors

Signed-off-by: Evgeny <[email protected]>

* make_quantizers() will take roles from the operation

Signed-off-by: Evgeny <[email protected]>

* Improve tests and fix missing imports

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply suggestions from code review

Signed-off-by: Tim Moon <[email protected]>

* Merge main followup

Signed-off-by: Evgeny <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Evgeny <[email protected]>
Signed-off-by: Evgeny Tsykunov <[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]>
@phu0ngng
[JAX] Rework amax reduction over TPSP (NVIDIA#2218)
ac4e0fd 
* rm using_global_amax_of_x

Signed-off-by: Phuong Nguyen <[email protected]>

---------

Signed-off-by: Phuong Nguyen <[email protected]>
@phu0ngng @jberchtold-nvidia
[JAX] Fix rng_state shape in fused attention (NVIDIA#2217)
b0d562d 
fix rng_state shape

Signed-off-by: Phuong Nguyen <[email protected]>
Co-authored-by: jberchtold-nvidia <[email protected]>
@negvet
[PyTorch] Fix QuantizedTensorBase -> QuantizedTensorStorage (NVIDIA#2226
ac886c3 
)

Fix QuantizedTensorBase -> QuantizedTensorStorage

Signed-off-by: Evgeny <[email protected]>
@ksivaman
Fix hang during debug build (NVIDIA#2221)
f0a9404 
Disable debug build for cutlass GEMM

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
115e528
@vthumbe1503
address review comments
d2072b1 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'gpt-oss-jax' of github.com:vthumbe1503/TransformerEngin…
6d9df80 
…e into gpt-oss-jax

Signed-off-by: Varun Thumbe <[email protected]>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
978fcde 
for more information, see https://pre-commit.ci
@vthumbe1503
minor bug because of merge conflict
13a3e3c 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'gpt-oss-jax' of github.com:vthumbe1503/TransformerEngin…
df0f449 
…e into gpt-oss-jax
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
d59526b 
for more information, see https://pre-commit.ci
@vthumbe1503
accept copilot suggestion
4783514 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'gpt-oss-jax' of github.com:vthumbe1503/TransformerEngin…
cda0c82 
…e into gpt-oss-jax
@ksivaman
Convert NVFP4BlockScaling to dataclass (NVIDIA#2227)
90449f7 
Fix passing args to nvfp4 recipe

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
@ptrendx @pre-commit-ci
Fix the cuBLAS workspace alignment (NVIDIA#2223)
aee5a82 
* Fix the cublas workspace alignment

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

* Fix

Co-authored-by: Copilot <[email protected]>
Signed-off-by: Przemyslaw Tredak <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Przemek Tredak <[email protected]>
Signed-off-by: Przemyslaw Tredak <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <[email protected]>
@vthumbe1503
fix test and remove a redundant test addition
14f8971 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
b9d7da7
@timmoon10
[PyTorch] Set usages for linear op quantizers before forward (NVIDIA#…
c100318 
…2222)

* Make sure to set usages for linear op quantizers before forward

Signed-off-by: Tim Moon <[email protected]>

* Avoid unsupported case for fused dbias+quantize kernel

Hopper does not support dbias + FP8 cast without FP8 transpose.

Signed-off-by: Tim Moon <[email protected]>

---------

Signed-off-by: Tim Moon <[email protected]>
@jberchtold-nvidia @phu0ngng
[JAX] Fix code block in fp8_autocast docstring (NVIDIA#2228)
f936c2a 
Fix code block in fp8_autocast docstring

Signed-off-by: Jeremy Berchtold <[email protected]>
Co-authored-by: Phuong Nguyen <[email protected]>
@jberchtold-nvidia @phu0ngng
[JAX] Fix shard map issue when get_all_mesh_axes() is used (NVIDIA#…
be7f43f 
…2229)

Fix shard map issue

Signed-off-by: Jeremy Berchtold <[email protected]>
Co-authored-by: Phuong Nguyen <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
6b0c73c
@hxbai @pre-commit-ci
@yaox12 @timmoon10
[PyTorch] fix int32 overflow in permute kernels (NVIDIA#2196)
e30c36a 
* fix overflow of int32 in permute kernels

Signed-off-by: Hongxiao Bai <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Hongxiao Bai <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Xin Yao <[email protected]>
Co-authored-by: Tim Moon <[email protected]>
@vthumbe1503
Merge branch 'main' into gpt-oss-jax
d462da1
@vthumbe1503
address review comments
5a55b0d 
Signed-off-by: Varun Thumbe <[email protected]>
@vthumbe1503
Merge branch 'gpt-oss-jax' of github.com:vthumbe1503/TransformerEngin…
bf3e04b 
…e into gpt-oss-jax

Signed-off-by: Varun Thumbe <[email protected]>
vthumbe1503 pushed a commit that referenced this pull request Jun 24, 2026
@cael-ling
Merge pull request #1 from vthumbe1503/feature/nvfp4-weight-swizzle-c…
8dafd75 
…ache

Enable weight swizzling for most cases
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.