Ifu dev 20260706 v2.17#660
Draft
AllenFarcas wants to merge 178 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>
…` with `total_recv_tokens_per_rank` placeholder (#3154) * versioning EP C configs Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Rename EP prepare token_counts to recv_tokens_per_expert Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Add total_recv_tokens_per_rank placeholder to nvte_ep_prepare Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Adapt PyTorch EP binding to versioned nvte_ep C config API Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Rename EP group config max_num_sms to num_comm_sms Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…L2 Jax dist (#3159) * Keep the routing map format alive Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix incorrectly launched multi process EP tests in L2 Jax instead of L2 jax dist Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
* disable 9.23.0/.1 for mxfp8 attention 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>
IFU of upstream release_v2.17 (156 commits, v2.15 -> v2.17) into ROCm dev. Resolved 78 conflicts and audited non-conflicting upstream changes for ROCm. Highlights: - NVFP4/MXFP4 enum chain: upstream enums keep values 8/9; ROCm-guarded MXFP4 enums renumbered to 10/11 across transformer_engine.h, common.h attr_sizes[], and transformer_engine.cpp switch cases. - _Linear refactored to upstream's LinearFwdArgs/LinearBwdArgs dataclass API; ROCm fields (keep_fp8_weight_transpose_cache, use_fsdp2, FSDP2 autocast) folded in. - Test Tensor helper adopted upstream Buffer/CudaPtr model; ROCm MXFP4 + gfx950 stochastic-rounding + NVTE_ROCM_BENCHMARK paths re-grafted. - Attention backend gating: re-fenced new NV-only cuDNN/SM disables behind not IS_HIP_EXTENSION; CP fused-attn API rename adopted with CK backend selection. - NCCL EP (new CUDA/NCCL-only subsystem) force-disabled on ROCm; topk moved to cuda_only (cub/cooperative_groups have no HIP equivalent) with jax guards; async_loader.h cuda_pipeline.h include guarded. Note: full ROCm HIP build + GPU runtime tests not yet run (no GPU/torch in the resolution environment). Static validation only: zero conflict markers, all Python py_compile clean, C++ preprocessor/brace balance verified, cross-file coherence audited. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
AllenFarcas
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
Post-merge fixes so the release_v2.17 IFU compiles and installs on ROCm
(gfx950): common library + PyTorch and JAX extensions.
Common library:
- quantize.cuh: guard CUDA-only nvfp4::compute_rowwise_amax (row-scaled NVFP4
unsupported on ROCm).
- fused_router/utils.h: restore sum()/max() helpers dropped by upstream's
warp_reduce_on_shmem rewrite; guard naive_topk_and_mask_smallk warp shuffles
for wave64 (non-sync width-32 variants).
- swizzle.cu: guard the new variable-shape grouped-swizzle kernel (wave32
reduction) on ROCm; keep the uniform-shape path.
- flash_attn.cu: guard the new TMA (CUtensorMap) multi_tensor_transpose_to_bhsd
path; ROCm stub for the entry point.
- rtc.cpp: guard the new CUDA-version diagnostic block.
- fused_router/async_loader.h: guard <cuda_pipeline.h>.
- fused_attn_rocm/fused_attn.cpp: reconcile nvte_fused_attn_fwd/bwd with the new
upstream format-descriptor params (strict on layout formats, ignore FP8
scale-inv formats).
- fused_attn.h: keep NVTE_F16_max512_seqlen on the CUDA path.
PyTorch extension:
- MXFP4Quantizer::create_tensor/create_grouped_tensor: match the new base
Quantizer virtuals (device/pin_memory; precomputed tensor_offsets).
JAX extension:
- Guard the new cuDNN score_mod (flex-attention) feature on ROCm; add a ROCm
stub for GetCudnnFrontendVersion.
- topk / EP handled elsewhere; NCCL EP forced off on ROCm.
Also: normalize ROCm guard comments ("Disabled on ROCm") and fix pre-existing
comment typos.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Upstream v2.17's Tensor test helper made scalar accessors (scale, amax, rowwise_scale_inv) strict: they assert the underlying buffer exists. Fork test call sites read them unconditionally, aborting on non-scaled tensors. Gate each call on the condition that governs buffer existence (is_fp8_output / isFp8Type input / MXFP4 has no global scale) so the accessors stay strict and still catch genuine misuse. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The v2.17 merge replaced the arch-aware fp8 e4m3 max with a hardcoded E4M3_MAX=448 in the dequant factor. On fnuz archs (e.g. gfx942/gfx950, maxNorm=240) this scales the dequantized output by 448/240 (~1.87x), producing wrong results. Restore the arch-aware factor on HIP device compilation using detail::TypeExtrema<fp8e4m3>::max; the CUDA path is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Align the upstream v2.17 operator tests with ROCm's actual support surface: - test_cast_nvfp4_transpose.cu: skip NVFP4 4over6 and row-scaled NVFP4 on ROCm (both are rejected by the production quantize guards) so they report SKIPPED instead of aborting on the assert. - test_multi_swizzle.cu: mirror the sibling swizzle suites. MultiTensorSwizzle runs on ROCm (skip gfx1250 pre-swizzle layout + shapes whose compact MXFP8 scales are not 128-aligned); Unswizzle and Roundtrip are compiled out, like UnswizzleTestSuite. - test_dequantize_nvfp4.cu: null the tensor amax so the dequant kernel takes the intended nullptr->1.0 fallback (the merged Tensor helper now allocates a zero-initialized amax buffer), matching the reference. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
inplace_swizzle_scale_for_gemm lacked the ROCm early-return that its sibling swizzle_scales_for_gemm has, so it wrote the nullopt swizzle result back and nulled valid NVFP4 scale_inv for optimize_for_gemm activations, crashing at generic_gemm. Mirror the sibling's guard (MXFP8/gfx1250 only). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The row-wise amax kernel was trapped in the Blackwell-only quantize_transpose_nvfp4 header (excluded from the ROCm build). Extract a portable nvfp4::compute_rowwise_amax into rowwise_amax_nvfp4.cuh and call it on both platforms; the quantize (vector_blockwise) and dequantize kernels already handle row-scaled inputs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The grouped-tensor GEMM entry points and fused grouped-linear path are CUDA-only; guard them with USE_ROCM / IS_HIP_EXTENSION so ROCm uses te_general_grouped_gemm. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move _QuantizeFunc / _IdentityFunc / _stride_from_shape imports to function scope to break the import cycle that prevented the module from loading. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Current scaling computes amax/scale inside the kernel, so supply workspace buffers instead of reading quantizer.scale/amax, use storage-safe size() and transpose-cache clears, and compare on _scale_inv in the test. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
RecipeState.create() passes roles= to every recipe-state class; MXFP4 was the only one missing it (parity with the other six). Update the MXFP4 test factory to dispatch on QuantizerRole fields instead of the obsolete string roles, and return a quantizer for every slot to satisfy make_quantizers. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Skip un-ported upstream features on ROCm: 4over6 and grouped NVFP4, and the Blackwell-only RHT cast-fusion swizzled-SF tests. Also fix tex->te.DType in test_permutation and add the bshd KV-cache FusedAttention ROCm skip. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
v2.17 gave recipes named ids (ids=recipe_id), so the positional 'recipe0' no longer matches and the mxfp8 test_numerics run selected 0 tests (pytest exit 5, counted as a harness error). Target the MXFP8 recipe by its name. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The fused quantized Triton-norm path on ROCm allocates the destination via make_empty with an internal quantizer, yielding an NVFP4TensorStorage; but update_quantized asserted the full NVFP4Tensor and aborted. Storage is the legitimate quantize target — the C++ fused-norm path writes into create_tensor output the same way, and tex.quantize operates on the storage buffers — so accept the NVFP4TensorStorage base type (NVFP4Tensor subclasses it, CUDA unaffected). Fixes NVFP4 activation-recompute under NVTE_USE_LAYERNORM_TRITON. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_FromNVFP4Func.forward moved a CPU-resident NVFP4 tensor to CUDA via the subclass .to(), which re-enters __torch_dispatch__ -> dequantize -> .to() and recurses infinitely (hit by FSDP2 DCP CPU staging). Move the raw buffers via the metadata dict instead (matching the _to_copy handler), avoiding the re-dispatch. Verified on mGPU: NVFP4 sync-DCP parity and the other NVFP4 fused-adam tests pass; recursion eliminated. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
fsdp_post_all_gather reconstructed the scale_inv via 1 / quantizer.scale, but Float8CurrentScalingQuantizer computes its scale in-kernel and has no .scale attribute, so current-scaling weights crashed with AttributeError during the FSDP2 all-gather (test_distributed[current-False]). fsdp_pre_all_gather already quantizes with amax reduction, so every rank's per-tensor scale_inv is identical; thread it through the all-gather metadata and use it for current scaling in post_all_gather. The delayed-scaling path (1 / quantizer.scale) is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
test_dcp_output_parity[NVFP4BlockScaling-async] is not bitwise-identical on torch 2.8: the default BlockingAsyncStager stages each parameter into a plain torch.empty(logical_shape, dtype) buffer, dequantizing the NVFP4Tensor. The checkpoint stores bf16 instead of the fp4 data + scales, and reload re-quantizes -- NVFP4's E4M3 block scales/amax do not survive the bf16 round-trip (MXFP8 survives via E8M0). Upstream #2721 fixes only Float8Tensor, via an aten.new_empty hook that torch 2.8's torch.empty-based stager never calls. Sync DCP is bitwise-correct, so checkpointing works. xfail is a temporary placeholder pending a proper fix (custom AsyncStager or torch upgrade). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
test_fp8_temp_accumulation_across_layers with quantized_model_init leaks ~0.68 MiB/layer for DelayedScaling/Float8CurrentScaling on ROCm: FSDP2's FSDPParam retains the all-gathered Float8Tensor reconstructed in the native fsdp_post_all_gather instead of freeing it on reshard (verified: the count of live reconstructed tensors grows unboundedly). This is FSDP2's unsharded-param free path for custom tensor subclasses, not the TE weight workspace -- skipping the workspace save was implemented and confirmed to have no effect. no_quant_init, sync/async DCP, and CUDA are unaffected. xfail is a temporary placeholder pending a proper fix. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
test_cpu_offloading_v1.py (multihead_attention, transformer_layer) failed with "ROCm fused attention: unsupported output/grad format" at NVTE_CPU_OFFLOAD_V1=1. When an interleaved QKV tensor is CPU-offloaded and reloaded, it comes back de-interleaved, so the backward sets ctx.qkv_layout to the de-interleaved reload layout (e.g. sbhd_sbhd_sbhd) while ctx.dqkv_layout stayed at the original interleaved layout (e.g. sb3hd). The ROCm fused-attn kernel uses a single qkv_layout for both inputs and gradients (it does not honor a separate dqkv_layout), so the nvte_fused_attn_bwd guard rejected the mismatch. On ROCm, make the grad layout follow the (possibly de-interleaved) input layout for the offload-reload case, matching pre-v2.17 behavior. Scoped to IS_HIP_EXTENSION and the reload path, so CUDA, non-offload, and MXFP8 layout handling are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
test_dpa_qkv_layout_thd[pad_between_seqs=True] failed with "No dot product attention backend is available" on ROCm. The test force-sets flash_attn_supported=True for the pad_between_seqs case (relying on _run_dot_product_attention to manually pad/unpad for the FlashAttention leg), but on ROCm the FlashAttention backend is disabled for THD with padding between sequences: flash-attn 2 cannot handle it and flash-attn 3 is not available. Forcing that leg left no usable backend and raised instead of comparing. Gate the override with `not IS_HIP_EXTENSION`. FusedAttention (CK) supports this case on ROCm, so the test now compares the backends that actually run. CUDA behavior is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The NVFP4 global encode scale hardcoded E4M3_MAX=448, but on fnuz archs (gfx942, maxNorm=240) the fp8 e4m3 scale storage and the dequant factor (fixed in 4e94fd7) use 240, so the encode scale was off by ~448/240, shifting quantized fp4 codes and failing the nvfp4 exact tests on gfx942. Use the arch-native fp8e4m3 max on HIP device compilation, mirroring the dequant fix. gfx950 (OCP, 448) is unchanged; CUDA path unchanged. NOTE: pending validation on gfx942 (rebuild + nvfp4 exact tests). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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: