Ifu dev 20260706 v2.17#660
Draft
AllenFarcas wants to merge 190 commits into
Draft
Conversation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Reduced padding kernel compilation time Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Completely removed unroll for better performance Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> --------- Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
…(#2905) Zero out padded region when swizzling via group quantize Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…2879) * fix broken links in README Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> * update README to modernize description and standardize terminology Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> * Update README.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> * Removed the duplicate line Signed-off-by: Przemek Tredak <ptredak@nvidia.com> --------- Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
* initial implementation for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * semi-working FP8; broken F16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * comment out F16 pass Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * pull in grouped_quantize for MXFP8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * grouped tensor - pytorch Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * quantize mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix shapes/strides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix unfused; clean up Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * split d to d_qk/d_v; attempt at bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at SWA/MLA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove leftover prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "update FE" This reverts commit d9ff5662aa4b4b6267c77baf614aada6602fa133. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix MLA O strides; add bottom_right_diagonal Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix get_quantizers; attempt at bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fprop; add o_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at bwd with o_format/d_out_format/dqkv_layout Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix dtype/o_format/etc in bwd calls Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix generateMatrixStridesWithFormats and _v1; fix padding for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix upon last commit for paddedsizes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add mxfp8 env var Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable FA for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add mha test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at bwd; force determinism; fix shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE from pre-merge branch to post-merge develop Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * allow MXFP8 linear + f16 attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * test cp a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints temporarily Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * test cp p2p Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes for mla Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * open up a2a for mla Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * test ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweaks for last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * enable mla ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert to main grouped tensor impl Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks to return to main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix combine_and_quantize for f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor tweaks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ds descale_o Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "fix ds descale_o" This reverts commit cd0bd82e239ff01210338b4e34cb8784109d22ec. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes for p2p and ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tweak cp test skips Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix bwd KV tensors Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak recipe control and backend selection Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak quantizer logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes after last two commits Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve generate strides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes for previous commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix bwd for current/delayed Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix dO/dO_f16 strides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix tests: SWA logic/test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add fp8 sink attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix a2a comm for F16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove nan/inf print in test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fa a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fa a2a+p2p f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to include new fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix thd for bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor a2a for fu/fa Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to fix d64 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor p2p/a2a+p2p; mostly regarding shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add shadow f16 fwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to fix SWA/BRCM Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * switch to GH FE temporarily Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * switch back to GL FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to latest commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update group tensor usage after merge main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * env vars for qdq(q,k), o_f16 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * allow other recipes than mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix grouped tensor for MLA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * change cp test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add shadow f16 bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix a2a+p2p for sbhd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix last commit and causal flag for fa Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * enable fp8 sink and disable fp8_mha Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor cleanup for cp/non-cp Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE for FP8 sink Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix TE for FP8 sink Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * temporary: random sink/print sink Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "temporary: random sink/print sink" This reverts commit 706095f802e04cbdd5d88ee53849cc5ec938203f. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace d_out_format with do_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix compare_and_assert for None cases Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove logic for b and simplify logic for dqkv types Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fix for ndim_q/kv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add explanation of fp8_output/grad in MHA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tidy up FP8 checks for bhsd/learnable Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove leading underscores in nvte_convert_qkv_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * simplify logic in generateMatrixStridesWithLayout Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up strides/ifelse-recipe logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak checks in utils.py Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak UnfusedDPA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * enable testing for ag+swa and disable fp8_mha Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak FusedAttn, fp8/f16 tensor naming/docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace d_out_format with do_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up p2p/a2a+p2p Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * qdq dO in bwd shadow f16 path Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak qdq dO logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints in shadow paths Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to allow non-determinism Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fuse qkv transposes; first pass Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remap parallelism to grid(bh, splits, 3) block(s/splits x d); use nvec = 128 bits Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * allocate contiguous block for qkv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix grouped tensor row/col scale_inv offsets Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * use fused permute kernels Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * quantize row/col as needed in fwd/bwd, non-cp/cp Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "quantize row/col as needed in fwd/bwd, non-cp/cp" This reverts commit ca5376956e8b8f662c7fa88661695b3e9eda4f8f. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Reapply "quantize row/col as needed in fwd/bwd, non-cp/cp" This reverts commit f19e852be3463210f2b3be5839ae8931e5ad92d0. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix v_col format when row is quantized Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add back necessary bwd quants for shadow paths/cp a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove ZInv for all layouts except T3HD Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix cp p2p with zinv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * temporarily switch to GH FE main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * switch back to GL FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix ag after merge main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add condition for qdq(do) to not affect other tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix custom_mha_fp8 test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix amax dqkv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fp8_recipe in DPA utils Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove use of amax for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add o_format/do_format/dqkv_layout to cache indicators for fp8 and f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * enable sink attn + FP8 in CP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to GH v1.22.0 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix for inconsistent kwarg name in permute to grouped tensor Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add TMA permute Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "add TMA permute" This reverts commit 2532a50e829144bee290fc94acb8f3f154a62ea9. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * TMA load for bhsd transposes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix some lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * temp: quant+perm+swizzle, rope, perm_fused Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove mla_rope for now; clean up quant+permute+pad_swizzle; create multi_tensor_swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * implement narrow-m for col swizzle; reorder to pad+perm+swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fused pad into perm; remove at::zeros as zeros done in perm kernels Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove shadow code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fix for permute shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * check smem size before entering narrow-k/m kernels Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * expand permute to multi_tensor_ Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor qkv/do quant; create a fast_path call Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [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 <8636796+cyanguwa@users.noreply.github.com> * cleanup grouped tensor fix Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove _with_amax for create_unquantized_tensor Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * reimplement inplace_multi_tensor_swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last commit; set swizzled flag in python Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove permute_to_grouped_tensor_bwd; clean up fwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add doxygen for multi_tensor_swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up nvte_convert_qkv_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes based on code review Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * group layouts/formats in APIs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename nvte_convert_qkv_format to nvte_convert_qkv_shape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove MXFP8 create_unquantized_tensor Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * rename permute_to_grouped_tensor to transpose_to_bhsd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add multi_tensor_swizzle_xx_unchecked and split the calls/paths Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * straighten up indexing for multi_tensor_pad Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * batch up kernel calls per-16-tensors for pad and permute Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove nvec128; rename nvec64 back to nvec Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add Macros/arch specifics for compilation Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [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 <8636796+cyanguwa@users.noreply.github.com> * attempt 1: MLA RoPE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "attempt 1: MLA RoPE" This reverts commit 79229248718d26a0ae7029206adc26c687bb42a7. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix kv_cache tests for Fused, is_page=True Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * attempt 2: MLA RoPE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * use DIVUP/_TO_MULTIPLE for pad_s_d_for_mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove CUDNN_VERSION 8900 macros Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add narrow-k/m swizzle tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * compile flash_attn.cu with special archs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "attempt 2: MLA RoPE" This reverts commit 3b854b29a3677de2005fecff821d801ccd9bf5d4. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * make contiguous instead of check is_contiguous Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unused s_q/s_kv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unused issue_tma_store_strided Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add version gate for mxfp8 for CPP users Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace nvte_get_qkv_shape with AttentionShape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * populate nvte_ changes to Jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to 1.22.1 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix minor merge issue Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert to FE 1.21 since it's what mxfp8 needs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * udpate jax attention shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert "revert to FE 1.21 since it's what mxfp8 needs" This reverts commit f09961a03bd7f5a316474b5d77b8292c7a49c1a6. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * pick FE 1.22 to support mxfp8 and avoid rng issue in 1.22.1 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix CP AG test on Hopper Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* starting grouped linear integration, not tested, grouped_bias_add optimized and uses scales now Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * all changes Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * optimize grouped bias add kernel to 4TB/s handling load imbalance Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * kernel optimized + review comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove uneecessary load from the kernel Remove unnecessary pointer reinterpretation for bias input. Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * No need for text for grouped bias add Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * Remove lru_cache from _is_deterministic_mode Removed lru_cache decorator from the _is_deterministic_mode function. Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * address more review comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * for better perf Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * remove unecessary comments Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename a bit Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add better ordering enforcment to split_overlap_rs gemms. This adds a short delay kernel to the split_overlap_rs function, which ensures that the gemms are properly ordered when run with cuda graphs. Signed-off-by: Chase Block <cblock@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Chase Block <cblock@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
…caches (#2805) * fixing mem leaks Signed-off-by: Peter St. John <pstjohn@nvidia.com> * update xfail message Signed-off-by: Peter St. John <pstjohn@nvidia.com> * addressing greptile comments Signed-off-by: Peter St. John <pstjohn@nvidia.com> * remove fsdp_safe Signed-off-by: Peter St. John <pstjohn@nvidia.com> * fix Float8BlockScaling backward override, unused imports, xfail MXFP8 HSDP Fix three issues: 1. LayerNormLinear weight quantizer was requesting columnwise usage even when backward_override is set. The FSDP2 refactor on this branch changed the columnwise condition to `is_grad_enabled and not is_fsdp2` but omitted the `backward_override is None` guard. This caused the weight to carry unnecessary columnwise data, which the dequantized backward path rejects. Also harden Float8BlockQuantizer::create_tensor to pass py::none() based on usage flags rather than relying on at::Tensor::defined() which is unreliable for default-constructed tensors in some PyTorch builds. 2. Remove unused imports: Float8Tensor from linear.py, and the entire float8_tensor import line from layernorm_linear.py (Float8Quantizer, Float8CurrentScalingQuantizer, Float8Tensor were all unused). 3. xfail MXFP8BlockScaling + fp8_init + HSDP in FSDP2 model tests. Pre-existing bug (confirmed on main) where fsdp_post_all_gather receives fewer output tensors than fsdp_pre_all_gather sent when the HSDP shard dimension is trivial (size 1). Signed-off-by: Peter St. John <pstjohn@nvidia.com> * address review comments Signed-off-by: Peter St. John <pstjohn@nvidia.com> * address greptile review Signed-off-by: Peter St. John <pstjohn@nvidia.com> * no need Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert grouped linear changes Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix all tests Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix cudnn commit Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix lint Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * revert Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix lint Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Peter St. John <pstjohn@nvidia.com> Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
…rentScalingQuantizer (#2816) * Make Float8CurrentScalingQuantizer stateless (no amax/scale members) Remove amax, scale, and use_existing_amax from Float8CurrentScalingQuantizer on both C++ and Python sides. All amax/scale allocations are now ad-hoc at quantization time: - quantize() allocates a combined 2-element tensor for amax+scale - quantize_with_amax() accepts amax as a parameter - create_unquantized_tensor_with_amax() returns amax in a tuple - set_quantization_params() is now a no-op Update all call sites in activation.cpp, bias.cpp, normalization.cpp, attention.cpp, and cast.cpp to propagate the amax buffer. For FusedAdam FP8 kernel: when scale_ptr is null (current scaling), derive scale from scale_inv and skip writing amax/scale_inv metadata. Python side passes empty(0) tensors for scale/amax to signal this. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard quantizer.amax/scale access with recipe.delayed() checks Add delayed() guards in context_parallel.py for amax writeback and scale cloning, since Float8CurrentScalingQuantizer no longer has these attributes. Allocate scratch scale buffer in _cast_master_weights_to_fp8_current_scaling instead of reading quantizer.scale. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * Fix pylint unused-argument warnings in Float8CurrentScalingQuantizer The device, use_existing_amax, scale, and amax parameters are kept for backward compatibility but not used internally. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * Clear dangling scale/amax pointers in Float8CurrentScalingQuantizer::quantize_impl amax_buf and scale_buf are caller-owned local tensors whose storage is released as soon as quantize() returns, leaving raw pointers stored in `out` dangling. No current caller reads out.scale()/out.amax() after quantize_impl returns, so this is currently safe, but it is a silent invariant that could turn into a use-after-free if new callers are added. Defensively clear both pointers at the end of quantize_impl (and in the empty-input early return), mirroring the existing set_amax(nullptr, ...) call already present before nvte_quantize_v2. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * Route current-scaling Float8Tensor through FP32 master path in FusedAdam Float8Tensor with Float8CurrentScalingQuantizer now goes through the FP32 master + requantize path (same as MXFP8/NVFP4/blockwise) instead of the fused FP8 Adam kernel. The fused FP8 kernel stays for delayed- scaling Float8Tensor only. Also revert adam.cu to upstream — current scaling no longer needs the scale_inv-derived path in the kernel. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore intent/perf comments in FusedAdam quantized dispatch - Note about a possible fused Adam+requantize kernel removing the FP32 round-trip for the QuantizedTensor path. - Justification for casting BF16/FP16 grads to FP32 before the FP32 Adam kernel. Also drop a stray f-string prefix on a literal-only RuntimeError part. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * tests skip Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * tests skip Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Warn on deprecated kwargs in Float8CurrentScalingQuantizer Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Guard remaining quantizer.scale/amax accesses with recipe.delayed() Followup to e06253d7. Three more spots accessed Float8CurrentScaling- Quantizer.scale/.amax (which no longer exist after the stateless refactor): - AttnFuncWithCPAndKVAllGather.forward used `not mxfp8()` instead of `delayed()` when cloning .scale, causing AttributeError under Float8CurrentScaling. - AttnFuncWithCPAndKVP2P fwd/bwd assigned .amax on per-step quantizer copies under `not mxfp8()`; harmless for current scaling (Python attaches a dynamic attribute that nobody reads) but inconsistent with the delayed()-guarded amax aggregation that follows. - print_quantizers (debug-only) read .scale/.amax for both DS and CS; restrict to DS only. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix fp8 and is_bwd_fp8 relationship Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
… (#2728) * fix: scope get_full_cu_seqlens cache key by device and inference mode Signed-off-by: Dongmin Ra <dongmin.ra@navercorp.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Dongmin Ra <dongmin.ra@navercorp.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Fix FA4 selection when FA3 is unavailable. Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
* Add AIR TopK support to TE JAX extension Adds a custom AIR TopK implementation (header-only, vendored into transformer_engine/common/util/) exposed as a JAX FFI custom call via the TE JAX extension. Key changes: - transformer_engine/common/util/air_topk.cu: AIR TopK CUDA kernel - transformer_engine/common/util/standalone_air_topk.cuh: vendored header - transformer_engine/common/include/transformer_engine/air_topk.h: C API - transformer_engine/jax/csrc/extensions/air_topk.cpp: JAX FFI binding - transformer_engine/jax/cpp_extensions/air_topk.py: Python wrapper - CMakeLists.txt: compile new kernel; use CCCL from CUDA toolkit - CMakeLists.txt: fix SM100 arch handling when all arches are special-cased Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Address PR review comments: fix namespace pollution, unused var, missing export, cache sm_cnt - Move WARP_SIZE/WARP_BITS/FULL_WARP_MASK/VECTORIZED_READ_SIZE into namespace nv - Remove unused keys_element_bytes variable in AirTopkFFI; collapse switch to dtype validation - Add missing `from .air_topk import *` export in jax/cpp_extensions/__init__.py - Cache sm_cnt per device with static vars to avoid repeated cudaGetDevice/cudaDeviceGetAttribute calls - Add CMAKE_BUILD_WITH_INSTALL_RPATH=ON to build_ext.py Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Rename air_topk -> topk throughout JAX extension Remove the `air_` prefix from all TopK-related identifiers: file names, C API functions (nvte_air_topk -> nvte_topk), FFI handler/primitive names (te_air_topk_ffi -> te_topk_ffi), Python symbols, and the internal `air_topk` namespace in standalone_topk.cuh. No functional changes. Signed-off-by: Diego Campora <dcampora@nvidia.com> Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Address ptrendx review comments and fix CI lint issues - Follow TE workspace convention: remove nvte_get_topk_workspace_bytes() and implement empty-workspace size-query pattern in nvte_topk() instead - Remove unnecessary nv_detail::float_to_T helper; replace usages with static_cast<T>() directly - Remove unrelated CMAKE_CUDA_ARCHITECTURES OFF block from CMakeLists.txt - Fix cpplint errors in standalone_topk.cuh: replace unsigned long long int with uint64_t, add NOLINT for constexpr-sized arrays and else-with-comment Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Add assertion for 2D input in topk Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --------- Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Signed-off-by: Diego Campora <dcampora@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…ed quantization kernel (#2921) Fix the race in the dbias computation Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* remove ctype to eliminate memory usage from the cudnn kernel Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * Remove c_dtype from fusible ops test Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Muu <koimuu@163.com>
* [Common, PyTorch] Add triton mHC kernels & pytorch operators Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * nit Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * make linter happy Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * nit Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * ah OK Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * new configs to improve perf Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * add APIs to docs Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix typos, check deterministic, refactor Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reset rng for all tests Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * add docstring Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix api doc Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * whoops Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * grad_x doesn't have to zero Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * nit Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * nit Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * force pytorch to not use bf16 for reduction Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * use TE's general_gemm instead Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Looks like this is how to make TE use fp32 acc Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…True (#2936) * fix Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add test Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * zero_out should also be tested Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root <root@gb-nvl-059-compute03.nvidia.com>
… (#2924) * Fix contiguous path for k=2880 Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * format Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Review suggestion from @Oleg-Goncharov Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Add test for swizzle + padding fusion Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Address review comments Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…942) * accumulate bias in fp32 instead of bf16 in ref impl dbias to avoid accumulated numerical error Signed-off-by: tdophung <tdophung@nvidia.com>
* loosen up thresholds. Also only check min loss of last 10% of steps to avoid failing by noise near convergence Signed-off-by: tdophung <tdophung@nvidia.com> * add deterministic flag for mnist run Signed-off-by: tdophung <tdophung@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* feat: add support for grouped GEMM swizzling with variable shapes and update C++ operator interface Signed-off-by: Abhishek <abhi.dtu11@gmail.com> * Added confirmation with uniformity in one of the dimensions Signed-off-by: Abhishek <abhi.dtu11@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Using single kernel for variable m and k Signed-off-by: Abhishek <abhi.dtu11@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cached blocks per sm for device and removed redundant checks Signed-off-by: Abhishek <abhi.dtu11@gmail.com> * Updated the code with newer changes in main Signed-off-by: Abhishek <abhi.dtu11@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Abhishek <abhi.dtu11@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
…929) * Avoid removing usages from quantized weight in linear op Quantized weight tensor may be used across steps, so removing a usage is not safe. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Tweak test to catch bug when alternating train and infer steps Signed-off-by: Tim Moon <tmoon@nvidia.com> * Avoid removing usages from quantized weights in grouped linear op Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore pre-forward quantizer config in ops Turns out we still need this in case the quantizer is used before the forward, e.g. in previous ops or CPU offloading. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Blindly preserve quantizer usages in quantized weight params. Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
… (#2943) * Patch for NVFP4 test suite Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * C++ tests fix Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Cleanup Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Removed dead code Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Set the golden value for amax in tests Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Fixed memory leakage Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> --------- Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
… (#2948) * Switch to cuDNN-FE min version 1.23.0 to enable fused grouped MLP Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix tests Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…n expert (#2947) Add workaround for cuteDSL stride requirement for zero token expert Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Remove unused NVFP4 quantize kernel Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Automated review summaryScope reviewed: ROCm-side commits on top of the NVIDIA v2.17 upstream merge (commits Code review: no additional correctness / safety findings beyond what @wenchenvincent has already flagged. The 21 existing inline questions on backend selection, ROCm-specific skips, fnuz handling, upstream API deltas, and dqkv layout still need author responses — this review does not duplicate them. Copyright headers: 11 files need updates — see inline comments. All are ROCm-modified/added files whose headers still carry only the NVIDIA line. Per repo policy, an AMD line at the current year (2026) should sit above the NVIDIA line, and new files added by ROCm need both headers. |
ipanfilo
requested changes
Jul 17, 2026
| @@ -158,24 +190,26 @@ def _run_layer_with_overlap( | |||
| raise AssertionError(result.stderr.decode()) | |||
|
|
|||
|
|
|||
| @pytest.mark.parametrize("use_cublasmp", (False, True)) | |||
Collaborator
There was a problem hiding this comment.
Should be only False for ROCm
Contributor
Author
There was a problem hiding this comment.
Gated use_cublasmp to (False,) if IS_HIP_EXTENSION else (False, True) on all six comm/gemm-overlap test parametrizations. On ROCm only the False cases are generated; CUDA still gets False+True.
| def test_distributed_fuser_ops(world_size: int) -> None: | ||
| """Launch parallel job that runs parallel tests""" | ||
| #TODO: find out why cannot align the following two lines with NV upstream | ||
| # TODO: find out why we cannot align the following two lines with NV upstream |
Collaborator
There was a problem hiding this comment.
TODO probably is not applicable - we know why the code is different. So it should rather be other comment, similar to test_cast_master_weights_to_fp8() which later solved the same problem
|
|
||
| # Swizzled SF is emitted only by the Blackwell RHT cast-fusion kernel; no fusion path on ROCm. | ||
| if IS_HIP_EXTENSION and optimize_for_gemm: | ||
| pytest.skip("RHT cast-fusion swizzled SF is Blackwell-only") |
Collaborator
There was a problem hiding this comment.
The message should rather say it is not supported on ROCm
Contributor
Author
There was a problem hiding this comment.
Corrected. Applied the same correction to the sibling module-level skip in test_nvfp4_rht_quantize_swizzle_fusion.py for consistency.
wangye805
requested changes
Jul 20, 2026
| if is_hip_extension(): | ||
| # Fully-masked striped steps must contribute nothing; the CK kernel can | ||
| # emit non-finite output/LSE for them, so neutralize to avoid 0 * NaN. | ||
| softmax_aux_per_step = jnp.where( |
Collaborator
There was a problem hiding this comment.
I think this issue is worth a new git ticket for fully understanding.
If it's a CK issue in gfx942, we may need to report to CK team.
Use pathlib.Path(sys.executable) / pathlib.Path(__file__).resolve() like upstream and test_cast_master_weights_to_fp8; drop the stale TODO. Verified test_distributed_fuser_ops passes on ROCm (world_size 1 and 2). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
NVFP4 on gfx942 was simulation-only and required fnuz-specific workarounds that were only needed because the fnuz e4m3 block scale (max 240) differs from OCP (448). Restrict NVFP4 to gfx950 (gfx1250 remains behind the existing TODO) and revert the fnuz workarounds: - quantization.py: NVFP4 arch gate ((9,4),(9,5)) -> ((9,5),); message updated. FP8 / block-scaling gates keep (9,4) (unaffected). - dequantize_nvfp4.cuh: drop the arch-aware factor_inv branch; use the upstream 1/(6*E4M3_MAX) form (no-op on gfx950/OCP). - quantization_ref_nvfp4.py: simplify all is_fp8_fnuz() branches to 448 and drop the now-unused import (no-op on gfx950). On gfx942 NVFP4 now reports unavailable, so all pytorch NVFP4 tests skip via the centralized _compute_nvfp4_support gate; gfx950 behavior is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Follow-up to bffa8f0, trimmed to the necessary changes: - quantization_ref_nvfp4.py: keep only the be3feea revert (global scale); restore the incidental is_fp8_fnuz cleanup (no-op on gfx950, unreached on gfx942) to keep the diff minimal. - test_dequantize_nvfp4.cu: skip on gfx942 (te_fp8_fnuz) — the dequant kernel now uses OCP 448 while its reference is fnuz-240, so it fails there. - test_nvfp4_rht_quantize_exact.py: gate the ROCm-only test_hadamard_transform_amax on recipe_available (was IS_HIP_EXTENSION-only; uses the reverted reference). Drops the unnecessary cast/grouped-gemm cpp skips (those tests pass on gfx942). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
68bc6e0 added the te_fp8_fnuz() skip to the CUDA-branch dequant bodies (#else of __HIP_PLATFORM_AMD__), which don't run on ROCm — so gfx942 still failed — and would break the CUDA build (te_fp8_fnuz is ROCm-only). Move the skip to the ROCm-branch DequantizeNVFP4 test body (the one that actually runs on gfx942) and drop it from the CUDA bodies. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@ipanfilo @wangye805 related to your comment on I cannot comment there but ill add here my reply: I replaced the TODO with an explanatory comment in the |
AllenFarcas
requested review from
ipanfilo,
matthiasdiener,
sudhu2k and
wangye805
July 22, 2026 19:18
wangye805
approved these changes
Jul 23, 2026
ipanfilo
reviewed
Jul 23, 2026
ipanfilo
left a comment
ipanfilo
left a comment
Collaborator
There was a problem hiding this comment.
Number of files have AMD copyright header incorrect format
Collaborator
There was a problem hiding this comment.
Look at pytorch.sh changes in release_v2.15_rocm, you might need the same here
| SM90_DET_FUSED_THD_BWD_MAX_BHSS = 1 << 30 | ||
| if ( | ||
| _deterministic | ||
| and qkv_format == "thd" |
Collaborator
There was a problem hiding this comment.
nit: add and not IS_HIP_EXTENSION
| python_exe = sys.executable | ||
| current_file = os.path.abspath(__file__) | ||
| python_exe = pathlib.Path(sys.executable) | ||
| current_file = pathlib.Path(__file__).resolve() |
Collaborator
There was a problem hiding this comment.
If it is different from upstream, add comment saying why
Contributor
Author
There was a problem hiding this comment.
The previous comment was a TODO to find out why we cannot allow with upstream, but right now that is solved and the change reflects upstream exactly.
Collaborator
There was a problem hiding this comment.
AMD copyright is missing
Collaborator
There was a problem hiding this comment.
The test needs to be added to pytorch.sh
| @@ -0,0 +1,119 @@ | |||
| /************************************************************************* | |||
| * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. | |||
Collaborator
There was a problem hiding this comment.
Missing header line about modificaton
ipanfilo
reviewed
Jul 23, 2026
| import transformer_engine_torch as tex | ||
|
|
||
| from transformer_engine.common.recipe import Recipe | ||
| from transformer_engine.pytorch.tensor._quantization_helpers import ( |
Collaborator
There was a problem hiding this comment.
So on ROCm _quantized_helpers import quntized_tensor? Why?
| @@ -96,9 +192,9 @@ def _compute_nvfp4_support() -> Tuple[bool, str]: | |||
| """Return if nvfp4 support is available""" | |||
| if IS_HIP_EXTENSION: | |||
| gpu_arch = get_device_compute_capability() | |||
| if gpu_arch in ((9, 4), (9, 5)): # TODO: enabled for gfx1250 when ready | |||
| if gpu_arch in ((9, 5),): # gfx950; TODO: enable for gfx1250 when ready | |||
Collaborator
There was a problem hiding this comment.
While it is expected to remove NVFP4 from gfx942, is there a reason to combine this change with IFU?
Contributor
Author
There was a problem hiding this comment.
gfx942 NVFP4 was simulation-only (fnuz E4M3 max is 240/256, not 448) and the v2.17 NVFP4 rework broke it. I tried fnuz workarounds to keep it alive, but for simplicity it is better to remove gfx942 NVFP4 here.
Contributor
Author
|
@ipanfilo In response to your question about The cycle is ROCm-specific because |
Contributor
Author
|
@ipanfilo I fixed all remaining AMD copyright and added any missing ones. |
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
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: