- 02 Jun, 2025 1 commit
-
-
jberchtold-nvidia authored
* Use 1x quantization + jax transpose on BW for performance Signed-off-by:
Jeremy Berchtold <jberchtold@nvidia.com> * Use 1x quantization on Hopper as well as it is also faster Signed-off-by:
-
- 29 May, 2025 3 commits
-
-
Hua Huang authored
* Support SWA in CP Ring Attn THD striped sharding Signed-off-by:
Hua Huang <huah@nvidia.com> * Add some comments; move check to _FusedAttnCPWithP2PHelper.check_supported() Signed-off-by:
-
Przemyslaw Tredak authored
* Changed the Tensor allocation strategy Signed-off-by:
Przemek Tredak <ptredak@nvidia.com> * Fixes Signed-off-by:
pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
-
Phuong Nguyen authored
* removes unnecessary reshapes for FP8 GEMM * use nn.jax.scaled_matmul Signed-off-by:
Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by:
-
- 28 May, 2025 2 commits
-
-
Phuong Nguyen authored
* fp8 gemm with direct quant Signed-off-by:
-
jberchtold-nvidia authored
Update jax_scaled_masked_softmax to match TE kernel implementation Signed-off-by:
-
- 22 May, 2025 2 commits
-
-
jberchtold-nvidia authored
Make primitive names more granular for better disabling granularity Signed-off-by: -
Kirthi Shankar Sivamani authored
* Build support for cuda 13 Signed-off-by:
Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix build for cudnn 8.9*; cuda 12.1 Signed-off-by:
-
- 16 May, 2025 1 commit
-
-
jberchtold-nvidia authored
* [JAX] Update flax module param initialization to support logical partitioning axes Signed-off-by:
-
- 15 May, 2025 1 commit
-
-
Kirthi Shankar Sivamani authored
* Cleanup runtime library loading Signed-off-by:
-
- 12 May, 2025 1 commit
-
-
jberchtold-nvidia authored
This reverts commit 5bee81e2 . Signed-off-by:
-
- 06 May, 2025 1 commit
-
-
jberchtold-nvidia authored
* Fix L2 test_custom_call_compute.py L2 tests Signed-off-by:
-
- 05 May, 2025 2 commits
-
-
Phuong Nguyen authored
* removes unneccessary reshapes for FP8 GEMM * use nn.jax.scaled_matmul Signed-off-by:
-
jberchtold-nvidia authored
* Enforce input sharding of norm primitive does not shard hidden dim Signed-off-by:
-
- 01 May, 2025 1 commit
-
-
Phuong Nguyen authored
* exclude GroupedGemm APIs Signed-off-by:
-
- 30 Apr, 2025 1 commit
-
-
jberchtold-nvidia authored
Fix distributed layernorm test failure Signed-off-by:
-
- 29 Apr, 2025 1 commit
-
-
jberchtold-nvidia authored
* Update test_helper.py and add QuantizeConfig class for CurrentScaling Signed-off-by:
-
- 28 Apr, 2025 1 commit
-
-
Kirthi Shankar Sivamani authored
* Add support for nvidia cu* lib wheels Signed-off-by:
-
- 24 Apr, 2025 1 commit
-
-
jberchtold-nvidia authored
Introduce nvte_memset to provide a fill kernel that is faster than cudaMemsetAsync for small sizes (#1716) * nvte_memset fills single float value Signed-off-by:
-
- 22 Apr, 2025 1 commit
-
-
jberchtold-nvidia authored
* [JAX-Q] Single GPU current scaling for JAX Signed-off-by:
-
- 21 Apr, 2025 1 commit
-
-
jberchtold-nvidia authored
Check CuDNN version and apply unfused norm if below a version with the fix Signed-off-by:
-
- 18 Apr, 2025 2 commits
-
-
Kirthi Shankar Sivamani authored
* Move jaxx cuda kernels to core Signed-off-by:
-
Phuong Nguyen authored
rm pax/praxis Signed-off-by:
-
- 17 Apr, 2025 1 commit
-
-
jberchtold-nvidia authored
* Add a flag to support computing zero-centered gamma in weight dtype or compute dtype for CuDNN Signed-off-by:
-
- 16 Apr, 2025 1 commit
-
-
Kshitij Lakhani authored
* Add test cases for full coverage in jax/test_layer.py - causal and window size None - causal and window size default (-1,1) - no_mask and window size default (-1,1) - no_mask and window size default (2,2) - padding and window size None - padding_causal and window_size (2,2) Signed-off-by:
Kshitij Janardan Lakhani <klakhani@nvidia.com> * Correct the condition where padding_causal_mask was being mapped to scaled upper triangle Signed-off-by:
-
- 14 Apr, 2025 2 commits
-
-
Johannes Reifferscheid authored
* Add experimental Shardy support. Production use is not yet recommended. --------- Signed-off-by:Johannes Reifferscheid <jreiffers@nvidia.com>
-
Hua Huang authored
* New GroupedGemmPrimitive using variadic args * Remove squeeze() to reduce D2D memcpy * Revert to the list append fashion to simplify code --------- Signed-off-by:
-
- 09 Apr, 2025 1 commit
-
-
Phuong Nguyen authored
* scaling enum abstract * rm NVTE_ from ScalingMode names * rework scaling mode enum in grouped gemm * fix norm sharding --------- Signed-off-by:
-
- 07 Apr, 2025 1 commit
-
-
Phuong Nguyen authored
* rm no scaling enum Signed-off-by:
-
- 04 Apr, 2025 3 commits
-
-
Phuong Nguyen authored
* rename QuantizeAxis to QuantizeLayout, get_layout to get_data_layout, q_axis to q_layout * add fatten_axis option * added gated act to test encoder * sharding constraint fixes * fix padding when flattening first dim needs to be padded * update test sizes so that padding is tested * rm output sharding as it can be done in the flax module * sharding scale_inv for mxfp8 --------- Signed-off-by: -
jberchtold-nvidia authored
MXFP8 flax layer tests Signed-off-by: -
kwyss-nvidia authored
* Blockwise float8 quantizer and quantized tensor class. The classes are configurable for 128x128 blocksize and 1x128 blocksize via setting block_scaling_dim == 2,1 respectively. Scale tensors are stored in a format emenable for matrix multiplication, however the integration of matmul is deferred as a separate story. Fusions of quantization and DBIAS or activation functions are not yet implemented, and the dequantization is currently implemented in torch. Tests for quantization are included in C++ and pytorch layers, with exact comparison to reference quantizer behavior as well as an attempt to hit interesting branches through the API such as tensor creation in pytorch and CPP and dequantization of row and columnwise usage. Two CUDA kernels for quantization are included, and are direct ports of equivalents in the kitchen repository, where a subchannel recipe has been used for end to end training. Signed-off-by:
Keith Wyss <kwyss@nvidia.com> * Apply linting changes. Signed-off-by:
Tim Moon <tmoon@nvidia.com> * Update JAX scaling mode enum Signed-off-by:
Tim Moon <4406448+timmoon10@users.noreply.github.com>
-
- 01 Apr, 2025 2 commits
-
-
Phuong Nguyen authored
* expose NVTE_FP8_COLLECTION_NAME, update_collections, get_delayed_scaling --------- Signed-off-by: -
Phuong Nguyen authored
* refactor + mxfp8 * added grouped gemm * rename linear to dense * added cublas init phase for groupedGemm * relax the tol of test encoder multiprocessing mxfp8 by 0.001 Signed-off-by:
-
- 31 Mar, 2025 1 commit
-
-
Michael Goldfarb authored
Add fast path for causal masking with segment IDs. Signed-off-by:Michael Goldfarb <mgoldfarb@nvidia.com>
-
- 18 Mar, 2025 2 commits
-
-
Michael Goldfarb authored
* Fix softmax shape for THD format. Signed-off-by: -
Charlene Yang authored
* add paged attention; test_kv_cache_accuray and test_paged_attn pass Signed-off-by:
Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unnecessary change from last commit Signed-off-by:
Charlene Yang <charleney@nvidia.com> * WIP: non-paged Signed-off-by:
-
- 13 Mar, 2025 2 commits
-
-
Reese Wang authored
Make ffi compatible with jax 0.4 Signed-off-by:
Reese Wang <rewang@nvidia.com> Co-authored-by:
-
Tim Moon authored
* Explicitly use python3 and pip3 Signed-off-by:
-
- 12 Mar, 2025 1 commit
-
-
Reese Wang authored
Remove xla_ignore_channel_id check and ignore Scan loop warning in unit test Signed-off-by:
-
