Release 2.15#656
Open
VeeraRajasekhar wants to merge 23 commits into
Open
Conversation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
…(#2905) Zero out padded region when swizzling via group quantize Signed-off-by: Kirthi Shankar Sivamani <ksivamani@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>
Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
fix fp8 and is_bwd_fp8 relationship Signed-off-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>
…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>
* [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>
…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>
…n expert (#2947) Add workaround for cuteDSL stride requirement for zero token expert Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…pace for THD sequences (#2522) * Get seqlens and offsets in O(N) space instead of O(N*N) space Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re enable fast causal path Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix: seqoffsets calculation for THD Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> * Clean up code. Add new comments. Fix unecessary pasing of seg pos to the seqoffsets calculation API Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Optimize and fix the slow O(T*T) path for seqlens and seqoffsets calculation for THD non-cp and Cp p2p ring - Newer path is O(T*max_segments) per seq - Newer path works well with CP p2p ring Fix BRCM cross attn by routing to new slow path rather than fast causal path Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix lint failure Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kshitij Janardan Lakhani <klakhani@login-ptyche02.ptyche.clusters.nvidia.com> Co-authored-by: JAX Toolbox <jax@nvidia.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>
…942) * accumulate bias in fp32 instead of bf16 in ref impl dbias to avoid accumulated numerical error Signed-off-by: tdophung <tdophung@nvidia.com>
…(#2955) * Better documentation for single param and envvar guard Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * fix doc Signed-off-by: ksivamani <ksivamani@nvidia.com> * Fix test envvar Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: ksivamani <ksivamani@nvidia.com>
- Resolve 9 merge conflicts (1 submodule + 8 content), preserving ROCm guards (CK backend selection, IS_HIP_EXTENSION gating) and adopting upstream refactors (swizzle check_scale_inv_shapes / unchecked variants, fused_attn FP8 validation removal, assert_close test helpers). - Take upstream cudnn-frontend submodule pin (strictly ahead, unmodified). - Port fused_attn_rocm/fused_attn.cpp nvte_fused_attn_fwd/bwd to upstream's new format-aware C API (o_format/do_format/dqkv_layout/scale_inv formats). ROCm CK/AOTriton backends derive layouts from qkv_layout and reject FP8/MXFP8, so new params are accepted for parity but unused.
- flash_attn.cu: guard upstream's new TMA (Tensor Memory Accelerator) transpose path (CUtensorMap / cp.async.bulk / ptx::mbarrier_*) behind #ifndef __HIP_PLATFORM_AMD__; ROCm uses the vectorized/shared-memory fallback paths. Force the TMA dispatch off on ROCm (sm_arch returns >=100 there). - swizzle.cu: move get_max_dynamic_smem() out of the CUDA-only guard so the ROCm-compiled swizzle kernels that call it can link (hipify maps the smem attribute to hipDeviceAttributeSharedMemPerBlockOptin).
Contributor
There was a problem hiding this comment.
Pull request overview
Release cut for TransformerEngine 2.15.0, including updates across PyTorch/JAX attention paths, grouped MLP/linear FP8+MXFP8 behavior, and scale swizzle utilities, plus expanded CI/QA coverage.
Changes:
- Add/extend MXFP8-related attention support (new QKV formats/layouts, softmax offset grad plumbing, CP updates) across Python and C++/CUDA backends.
- Refine grouped linear/MLP behavior (env-gated single-param mode, cuDNN frontend gating, updated quantizer usage, wgrad path updates).
- Add multi-tensor scale swizzle helpers + tests; adjust build/QA scripts for new coverage and release versioning.
Reviewed changes
Copilot reviewed 61 out of 63 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| transformer_engine/pytorch/utils.py | Add env-gated resolver for grouped-linear single-param flags (with warnings). |
| transformer_engine/pytorch/triton/init.py | Export mhc from the Triton wrapper package. |
| transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py | Adjust scale shape/offset computations for quantized grouped tensors (incl. block scaling). |
| transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py | Update fused grouped-MLP forward gating and cuDNN kwargs handling. |
| transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py | Update fused grouped-MLP backward gating, wgrad kernel selection, and tensor packing edge cases. |
| transformer_engine/pytorch/ops/basic/grouped_linear.py | Gate single grouped weight/bias via env var; refine weight quantizer usage + recipe-state reset behavior. |
| transformer_engine/pytorch/ops/basic/basic_linear.py | Refine weight quantizer usage + recipe-state reset behavior for BasicLinear. |
| transformer_engine/pytorch/ops/_common.py | Consolidate cuDNN frontend version gate for fused grouped-MLP support. |
| transformer_engine/pytorch/newton_schulz.py | Change coefficient API to structured tuples and flatten only at C++ call boundary. |
| transformer_engine/pytorch/module/grouped_linear.py | Gate single grouped weight/bias via env var in module wrapper. |
| transformer_engine/pytorch/csrc/util.h | Add unchecked multi-tensor swizzle API declaration. |
| transformer_engine/pytorch/csrc/extensions/swizzle.cpp | Implement unchecked swizzle path and add pybind-facing multi-tensor swizzle helpers; improve grouped swizzle padding handling. |
| transformer_engine/pytorch/csrc/extensions/pybind.cpp | Bind multi-tensor swizzle + fused-attn tensor utilities into the PyTorch extension. |
| transformer_engine/pytorch/csrc/extensions.h | Extend fused attention signatures and declare multi-tensor transpose/pad APIs. |
| transformer_engine/pytorch/cpp_extensions/fused_attn.py | Extend fused-attn Python wrapper API with output/scale formats and richer backward outputs. |
| transformer_engine/pytorch/attention/multi_head_attention.py | Adjust FP8 output decisions to account for MXFP8 scaling constraints. |
| transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py | Extend DPA recipe routing for MXFP8; update layout detection and CP call wiring. |
| transformer_engine/jax/csrc/extensions/attention.cpp | Fix ragged vs non-ragged fused-attn tensor shapes and pass new format args through. |
| transformer_engine/jax/attention.py | Replace O(T²) mask path with O(T·S) per-segment aggregation path for THD seqlen/offset metadata. |
| transformer_engine/common/util/pybind_helper.h | Add new QKV format/layout enum values to Python bindings. |
| transformer_engine/common/transformer_engine.cpp | Make scale shape checking optionally skippable in input validation. |
| transformer_engine/common/include/transformer_engine/swizzle.h | Add unchecked multi-tensor swizzle API for cases with intentionally mismatched scale shapes. |
| transformer_engine/common/include/transformer_engine/fused_attn.h | Extend fused-attn API with new formats/layouts and add a C++ AttentionShape helper. |
| transformer_engine/common/fused_attn/utils.h | Add MXFP8 padding helpers, stride generation by format, and extend fused-attn descriptor keys. |
| transformer_engine/common/fused_attn/utils.cu | Add stride generation support for BHSD layout. |
| transformer_engine/common/fused_attn/fused_attn.cpp | Plumb new fused-attn format/layout args; extend backend eligibility for BHSD and newer cuDNN versions. |
| transformer_engine/common/fused_attn/fused_attn_fp8.h | Update FP8 fused-attn forward/backward signatures to match new format/layout plumbing. |
| transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h | Remove cuDNN version compile-guard wrappers from header. |
| transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu | Remove cuDNN version compile-guard wrappers from compilation unit. |
| transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h | Extend f16 arbitrary-seqlen signatures with output/grad formats and dqkv layout. |
| transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu | Plumb new format/layout args through arbitrary-seqlen kernels and caching keys. |
| transformer_engine/common/fused_attn_rocm/fused_attn.cpp | Accept (but ignore) new fused-attn format/layout args on ROCm for API parity. |
| transformer_engine/common/common.h | Add integer vectorized-load macro and update CheckInputTensor signature default arg. |
| transformer_engine/common/CMakeLists.txt | Move flash_attn.cu into arch-specific CUDA source list. |
| transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh | Add synchronization before dbias reuse path. |
| transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh | Fix scale indexing for swizzled scales + zero OOB padding; add sync before dbias reuse path. |
| tests/pytorch/utils.py | Handle None outputs in compare helper. |
| tests/pytorch/distributed/run_newton_schulz.py | Update reference to new coefficient tuple format. |
| tests/pytorch/attention/test_attention.py | Expand attention test matrix (new dims, new scaling mode, softmax offset grad, new formats). |
| tests/pytorch/attention/test_attention_with_cp.py | Expand CP attention test coverage (new configs, determinism plumbing, MXFP8 scaling mode). |
| tests/pytorch/attention/run_attention_with_cp.py | Add MXFP8 quantizers + determinism env plumbing; improve NaN/Inf asserts. |
| tests/jax/test_custom_call_compute.py | Make grouped-dense bias grad reference accumulate in fp32 to match primitive. |
| tests/cpp/operator/test_swizzle.cu | Add compact grouped-scale swizzle test and broaden narrow-path coverage. |
| tests/cpp/operator/test_multi_swizzle.cu | Add new C++ tests for multi-tensor swizzle/unswizzle/roundtrip. |
| tests/cpp/operator/CMakeLists.txt | Include new multi-swizzle test; disable it on ROCm for now. |
| qa/L0_pytorch_unittest/test.sh | Enable grouped-linear single-param mode for relevant unit tests; add MHC unittest invocation. |
| qa/L0_pytorch_lint/test.sh | Add a lint QA job for cpplint + pylint. |
| qa/L0_pytorch_debug_unittest/test.sh | Enable grouped-linear single-param mode for debug unittests. |
| docs/api/pytorch.rst | Add API docs entries for Triton MHC fused functions. |
| ci/pytorch.sh | Enable grouped-linear single-param mode for fusible-ops and sanity tests in CI. |
| build_tools/VERSION.txt | Bump version from 2.15.0.dev0 to 2.15.0. |
Comments suppressed due to low confidence (1)
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py:317
grouped_gemm_wgrad_kernel()is annotated/used as an optional kernel selector, but it can currently raiseImportErrorwhengrouped_gemm_wgrad_wrapper_sm100is unavailable. Since callers passcudnn_wgrad_kernel_fn=self.grouped_gemm_wgrad_kernel()unconditionally, this can crash at runtime even though a safe fallback path exists (i.e., returningNoneso the non-cuDNN wgrad path is used). Catch the import failure and returnNoneinstead.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # See LICENSE for license information. | ||
|
|
||
| """PyTorch wrappers for Triton kernels.""" | ||
| from transformer_engine.pytorch.triton import mhc |
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
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: