Ifu dev 20260706 v2.17#660
Conversation
Signed-off-by: Przemek Tredak <[email protected]>
* Reduced padding kernel compilation time Signed-off-by: Oleg Goncharov <[email protected]> * Completely removed unroll for better performance Signed-off-by: Oleg Goncharov <[email protected]> --------- Signed-off-by: Oleg Goncharov <[email protected]>
…(#2905) Zero out padded region when swizzling via group quantize Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
…2879) * fix broken links in README Signed-off-by: Santosh Bhavani <[email protected]> * update README to modernize description and standardize terminology Signed-off-by: Santosh Bhavani <[email protected]> * Update README.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Santosh Bhavani <[email protected]> * Removed the duplicate line Signed-off-by: Przemek Tredak <[email protected]> --------- Signed-off-by: Santosh Bhavani <[email protected]> Signed-off-by: Przemek Tredak <[email protected]> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Przemek Tredak <[email protected]>
* initial implementation for mxfp8 Signed-off-by: Charlene Yang <[email protected]> * semi-working FP8; broken F16 Signed-off-by: Charlene Yang <[email protected]> * clean up last commit Signed-off-by: Charlene Yang <[email protected]> * comment out F16 pass Signed-off-by: Charlene Yang <[email protected]> * pull in grouped_quantize for MXFP8 Signed-off-by: Charlene Yang <[email protected]> * grouped tensor - pytorch Signed-off-by: Charlene Yang <[email protected]> * quantize mxfp8 Signed-off-by: Charlene Yang <[email protected]> * fix shapes/strides Signed-off-by: Charlene Yang <[email protected]> * fix unfused; clean up Signed-off-by: Charlene Yang <[email protected]> * split d to d_qk/d_v; attempt at bwd Signed-off-by: Charlene Yang <[email protected]> * fix last merge Signed-off-by: Charlene Yang <[email protected]> * update FE Signed-off-by: Charlene Yang <[email protected]> * attempt at SWA/MLA Signed-off-by: Charlene Yang <[email protected]> * remove prints Signed-off-by: Charlene Yang <[email protected]> * remove leftover prints Signed-off-by: Charlene Yang <[email protected]> * Revert "update FE" This reverts commit d9ff5662aa4b4b6267c77baf614aada6602fa133. Signed-off-by: Charlene Yang <[email protected]> * update FE Signed-off-by: Charlene Yang <[email protected]> * fix MLA O strides; add bottom_right_diagonal Signed-off-by: Charlene Yang <[email protected]> * attempt at bwd Signed-off-by: Charlene Yang <[email protected]> * fix get_quantizers; attempt at bwd Signed-off-by: Charlene Yang <[email protected]> * fix fprop; add o_format Signed-off-by: Charlene Yang <[email protected]> * attempt at bwd with o_format/d_out_format/dqkv_layout Signed-off-by: Charlene Yang <[email protected]> * fix dtype/o_format/etc in bwd calls Signed-off-by: Charlene Yang <[email protected]> * fix generateMatrixStridesWithFormats and _v1; fix padding for mxfp8 Signed-off-by: Charlene Yang <[email protected]> * fix upon last commit for paddedsizes Signed-off-by: Charlene Yang <[email protected]> * add mxfp8 env var Signed-off-by: Charlene Yang <[email protected]> * disable FA for mxfp8 Signed-off-by: Charlene Yang <[email protected]> * add mha test Signed-off-by: Charlene Yang <[email protected]> * attempt at bwd; force determinism; fix shapes Signed-off-by: Charlene Yang <[email protected]> * remove prints Signed-off-by: Charlene Yang <[email protected]> * update FE Signed-off-by: Charlene Yang <[email protected]> * update FE from pre-merge branch to post-merge develop Signed-off-by: Charlene Yang <[email protected]> * allow MXFP8 linear + f16 attn Signed-off-by: Charlene Yang <[email protected]> * test cp a2a Signed-off-by: Charlene Yang <[email protected]> * remove prints temporarily Signed-off-by: Charlene Yang <[email protected]> * test cp p2p Signed-off-by: Charlene Yang <[email protected]> * minor fixes for mla Signed-off-by: Charlene Yang <[email protected]> * open up a2a for mla Signed-off-by: Charlene Yang <[email protected]> * test ag Signed-off-by: Charlene Yang <[email protected]> * tweaks for last commit Signed-off-by: Charlene Yang <[email protected]> * enable mla ag Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge Signed-off-by: Charlene Yang <[email protected]> * revert to main grouped tensor impl Signed-off-by: Charlene Yang <[email protected]> * minor tweaks to return to main Signed-off-by: Charlene Yang <[email protected]> * remove prints Signed-off-by: Charlene Yang <[email protected]> * fix combine_and_quantize for f16 Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor tweaks Signed-off-by: Charlene Yang <[email protected]> * tweak tests Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ds descale_o Signed-off-by: Charlene Yang <[email protected]> * Revert "fix ds descale_o" This reverts commit cd0bd82e239ff01210338b4e34cb8784109d22ec. Signed-off-by: Charlene Yang <[email protected]> * minor fixes for p2p and ag Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tweak cp test skips Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE Signed-off-by: Charlene Yang <[email protected]> * fix bwd KV tensors Signed-off-by: Charlene Yang <[email protected]> * tweak recipe control and backend selection Signed-off-by: Charlene Yang <[email protected]> * tweak quantizer logic Signed-off-by: Charlene Yang <[email protected]> * minor fixes after last two commits Signed-off-by: Charlene Yang <[email protected]> * improve generate strides Signed-off-by: Charlene Yang <[email protected]> * minor fixes for previous commit Signed-off-by: Charlene Yang <[email protected]> * fix bwd for current/delayed Signed-off-by: Charlene Yang <[email protected]> * tweak test configs Signed-off-by: Charlene Yang <[email protected]> * fix dO/dO_f16 strides Signed-off-by: Charlene Yang <[email protected]> * fix tests: SWA logic/test configs Signed-off-by: Charlene Yang <[email protected]> * fix ag Signed-off-by: Charlene Yang <[email protected]> * add fp8 sink attn Signed-off-by: Charlene Yang <[email protected]> * fix a2a comm for F16 Signed-off-by: Charlene Yang <[email protected]> * remove nan/inf print in test Signed-off-by: Charlene Yang <[email protected]> * fix fa a2a Signed-off-by: Charlene Yang <[email protected]> * fix fa a2a+p2p f16 Signed-off-by: Charlene Yang <[email protected]> * update FE to include new fixes Signed-off-by: Charlene Yang <[email protected]> * fix thd for bwd Signed-off-by: Charlene Yang <[email protected]> * refactor a2a for fu/fa Signed-off-by: Charlene Yang <[email protected]> * update FE to fix d64 Signed-off-by: Charlene Yang <[email protected]> * refactor ag Signed-off-by: Charlene Yang <[email protected]> * refactor p2p/a2a+p2p; mostly regarding shapes Signed-off-by: Charlene Yang <[email protected]> * add shadow f16 fwd Signed-off-by: Charlene Yang <[email protected]> * update FE to fix SWA/BRCM Signed-off-by: Charlene Yang <[email protected]> * switch to GH FE temporarily Signed-off-by: Charlene Yang <[email protected]> * switch back to GL FE Signed-off-by: Charlene Yang <[email protected]> * update FE to latest commit Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update group tensor usage after merge main Signed-off-by: Charlene Yang <[email protected]> * env vars for qdq(q,k), o_f16 tests Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * allow other recipes than mxfp8 Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix grouped tensor for MLA Signed-off-by: Charlene Yang <[email protected]> * change cp test configs Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add shadow f16 bwd Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix a2a+p2p for sbhd Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix last commit and causal flag for fa Signed-off-by: Charlene Yang <[email protected]> * enable fp8 sink and disable fp8_mha Signed-off-by: Charlene Yang <[email protected]> * minor cleanup for cp/non-cp Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE for FP8 sink Signed-off-by: Charlene Yang <[email protected]> * fix TE for FP8 sink Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * temporary: random sink/print sink Signed-off-by: Charlene Yang <[email protected]> * Revert "temporary: random sink/print sink" This reverts commit 706095f802e04cbdd5d88ee53849cc5ec938203f. Signed-off-by: Charlene Yang <[email protected]> * replace d_out_format with do_format Signed-off-by: Charlene Yang <[email protected]> * fix compare_and_assert for None cases Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove logic for b and simplify logic for dqkv types Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fix for ndim_q/kv Signed-off-by: Charlene Yang <[email protected]> * add explanation of fp8_output/grad in MHA Signed-off-by: Charlene Yang <[email protected]> * tidy up FP8 checks for bhsd/learnable Signed-off-by: Charlene Yang <[email protected]> * remove leading underscores in nvte_convert_qkv_format Signed-off-by: Charlene Yang <[email protected]> * simplify logic in generateMatrixStridesWithLayout Signed-off-by: Charlene Yang <[email protected]> * clean up strides/ifelse-recipe logic Signed-off-by: Charlene Yang <[email protected]> * tweak checks in utils.py Signed-off-by: Charlene Yang <[email protected]> * tweak UnfusedDPA Signed-off-by: Charlene Yang <[email protected]> * enable testing for ag+swa and disable fp8_mha Signed-off-by: Charlene Yang <[email protected]> * tweak FusedAttn, fp8/f16 tensor naming/docstring Signed-off-by: Charlene Yang <[email protected]> * replace d_out_format with do_format Signed-off-by: Charlene Yang <[email protected]> * fix lint Signed-off-by: Charlene Yang <[email protected]> * clean up a2a Signed-off-by: Charlene Yang <[email protected]> * clean up ag Signed-off-by: Charlene Yang <[email protected]> * clean up p2p/a2a+p2p Signed-off-by: Charlene Yang <[email protected]> * tweak test configs Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * qdq dO in bwd shadow f16 path Signed-off-by: Charlene Yang <[email protected]> * tweak qdq dO logic Signed-off-by: Charlene Yang <[email protected]> * remove prints in shadow paths Signed-off-by: Charlene Yang <[email protected]> * update FE to allow non-determinism Signed-off-by: Charlene Yang <[email protected]> * fuse qkv transposes; first pass Signed-off-by: Charlene Yang <[email protected]> * remap parallelism to grid(bh, splits, 3) block(s/splits x d); use nvec = 128 bits Signed-off-by: Charlene Yang <[email protected]> * allocate contiguous block for qkv Signed-off-by: Charlene Yang <[email protected]> * fix grouped tensor row/col scale_inv offsets Signed-off-by: Charlene Yang <[email protected]> * use fused permute kernels Signed-off-by: Charlene Yang <[email protected]> * quantize row/col as needed in fwd/bwd, non-cp/cp Signed-off-by: Charlene Yang <[email protected]> * Revert "quantize row/col as needed in fwd/bwd, non-cp/cp" This reverts commit ca5376956e8b8f662c7fa88661695b3e9eda4f8f. Signed-off-by: Charlene Yang <[email protected]> * Reapply "quantize row/col as needed in fwd/bwd, non-cp/cp" This reverts commit f19e852be3463210f2b3be5839ae8931e5ad92d0. Signed-off-by: Charlene Yang <[email protected]> * fix v_col format when row is quantized Signed-off-by: Charlene Yang <[email protected]> * add back necessary bwd quants for shadow paths/cp a2a Signed-off-by: Charlene Yang <[email protected]> * remove ZInv for all layouts except T3HD Signed-off-by: Charlene Yang <[email protected]> * fix cp p2p with zinv Signed-off-by: Charlene Yang <[email protected]> * temporarily switch to GH FE main Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * switch back to GL FE Signed-off-by: Charlene Yang <[email protected]> * fix ag after merge main Signed-off-by: Charlene Yang <[email protected]> * add condition for qdq(do) to not affect other tests Signed-off-by: Charlene Yang <[email protected]> * fix custom_mha_fp8 test Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix amax dqkv Signed-off-by: Charlene Yang <[email protected]> * fix fp8_recipe in DPA utils Signed-off-by: Charlene Yang <[email protected]> * remove use of amax for mxfp8 Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add o_format/do_format/dqkv_layout to cache indicators for fp8 and f16 Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * enable sink attn + FP8 in CP Signed-off-by: Charlene Yang <[email protected]> * update FE to GH v1.22.0 Signed-off-by: Charlene Yang <[email protected]> * fix for inconsistent kwarg name in permute to grouped tensor Signed-off-by: Charlene Yang <[email protected]> * add TMA permute Signed-off-by: Charlene Yang <[email protected]> * Revert "add TMA permute" This reverts commit 2532a50e829144bee290fc94acb8f3f154a62ea9. Signed-off-by: Charlene Yang <[email protected]> * TMA load for bhsd transposes Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix some lint Signed-off-by: Charlene Yang <[email protected]> * temp: quant+perm+swizzle, rope, perm_fused Signed-off-by: Charlene Yang <[email protected]> * remove mla_rope for now; clean up quant+permute+pad_swizzle; create multi_tensor_swizzle Signed-off-by: Charlene Yang <[email protected]> * fix last commit Signed-off-by: Charlene Yang <[email protected]> * implement narrow-m for col swizzle; reorder to pad+perm+swizzle Signed-off-by: Charlene Yang <[email protected]> * fused pad into perm; remove at::zeros as zeros done in perm kernels Signed-off-by: Charlene Yang <[email protected]> * remove shadow code Signed-off-by: Charlene Yang <[email protected]> * minor fix for permute shapes Signed-off-by: Charlene Yang <[email protected]> * check smem size before entering narrow-k/m kernels Signed-off-by: Charlene Yang <[email protected]> * expand permute to multi_tensor_ Signed-off-by: Charlene Yang <[email protected]> * refactor qkv/do quant; create a fast_path call Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Charlene Yang <[email protected]> * cleanup grouped tensor fix Signed-off-by: Charlene Yang <[email protected]> * remove _with_amax for create_unquantized_tensor Signed-off-by: Charlene Yang <[email protected]> * fix last commit Signed-off-by: Charlene Yang <[email protected]> * reimplement inplace_multi_tensor_swizzle Signed-off-by: Charlene Yang <[email protected]> * fix last commit; set swizzled flag in python Signed-off-by: Charlene Yang <[email protected]> * remove permute_to_grouped_tensor_bwd; clean up fwd Signed-off-by: Charlene Yang <[email protected]> * add doxygen for multi_tensor_swizzle Signed-off-by: Charlene Yang <[email protected]> * clean up nvte_convert_qkv_format Signed-off-by: Charlene Yang <[email protected]> * fixes based on code review Signed-off-by: Charlene Yang <[email protected]> * group layouts/formats in APIs Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename nvte_convert_qkv_format to nvte_convert_qkv_shape Signed-off-by: Charlene Yang <[email protected]> * remove MXFP8 create_unquantized_tensor Signed-off-by: Charlene Yang <[email protected]> * rename permute_to_grouped_tensor to transpose_to_bhsd Signed-off-by: Charlene Yang <[email protected]> * add multi_tensor_swizzle_xx_unchecked and split the calls/paths Signed-off-by: Charlene Yang <[email protected]> * straighten up indexing for multi_tensor_pad Signed-off-by: Charlene Yang <[email protected]> * batch up kernel calls per-16-tensors for pad and permute Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove nvec128; rename nvec64 back to nvec Signed-off-by: Charlene Yang <[email protected]> * add Macros/arch specifics for compilation Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Charlene Yang <[email protected]> * attempt 1: MLA RoPE Signed-off-by: Charlene Yang <[email protected]> * Revert "attempt 1: MLA RoPE" This reverts commit 79229248718d26a0ae7029206adc26c687bb42a7. Signed-off-by: Charlene Yang <[email protected]> * fix kv_cache tests for Fused, is_page=True Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * attempt 2: MLA RoPE Signed-off-by: Charlene Yang <[email protected]> * use DIVUP/_TO_MULTIPLE for pad_s_d_for_mxfp8 Signed-off-by: Charlene Yang <[email protected]> * remove CUDNN_VERSION 8900 macros Signed-off-by: Charlene Yang <[email protected]> * add narrow-k/m swizzle tests Signed-off-by: Charlene Yang <[email protected]> * compile flash_attn.cu with special archs Signed-off-by: Charlene Yang <[email protected]> * Revert "attempt 2: MLA RoPE" This reverts commit 3b854b29a3677de2005fecff821d801ccd9bf5d4. Signed-off-by: Charlene Yang <[email protected]> * make contiguous instead of check is_contiguous Signed-off-by: Charlene Yang <[email protected]> * remove unused s_q/s_kv Signed-off-by: Charlene Yang <[email protected]> * remove unused issue_tma_store_strided Signed-off-by: Charlene Yang <[email protected]> * add version gate for mxfp8 for CPP users Signed-off-by: Charlene Yang <[email protected]> * replace nvte_get_qkv_shape with AttentionShape Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * populate nvte_ changes to Jax Signed-off-by: Charlene Yang <[email protected]> * update FE to 1.22.1 Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix minor merge issue Signed-off-by: Charlene Yang <[email protected]> * revert to FE 1.21 since it's what mxfp8 needs Signed-off-by: Charlene Yang <[email protected]> * udpate jax attention shapes Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert "revert to FE 1.21 since it's what mxfp8 needs" This reverts commit f09961a03bd7f5a316474b5d77b8292c7a49c1a6. Signed-off-by: Charlene Yang <[email protected]> * pick FE 1.22 to support mxfp8 and avoid rng issue in 1.22.1 Signed-off-by: Charlene Yang <[email protected]> * fix CP AG test on Hopper Signed-off-by: Charlene Yang <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Charlene Yang <[email protected]> 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 <[email protected]> * all changes Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * [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 <[email protected]> * [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 <[email protected]> * No need for text for grouped bias add Signed-off-by: vthumbe1503 <[email protected]> * Remove lru_cache from _is_deterministic_mode Removed lru_cache decorator from the _is_deterministic_mode function. Signed-off-by: vthumbe1503 <[email protected]> * address more review comments Signed-off-by: Varun Thumbe <[email protected]> * for better perf Signed-off-by: Varun Thumbe <[email protected]> * remove unecessary comments Signed-off-by: vthumbe1503 <[email protected]> * [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 <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> 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 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Chase Block <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak <[email protected]>
Signed-off-by: Björn Buschkämper <[email protected]>
…caches (#2805) * fixing mem leaks Signed-off-by: Peter St. John <[email protected]> * update xfail message Signed-off-by: Peter St. John <[email protected]> * addressing greptile comments Signed-off-by: Peter St. John <[email protected]> * remove fsdp_safe Signed-off-by: Peter St. John <[email protected]> * 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 <[email protected]> * address review comments Signed-off-by: Peter St. John <[email protected]> * address greptile review Signed-off-by: Peter St. John <[email protected]> * no need Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * fix all tests Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * fix lint Signed-off-by: Varun Thumbe <[email protected]> * revert Signed-off-by: Varun Thumbe <[email protected]> * fix lint Signed-off-by: Varun Thumbe <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Peter St. John <[email protected]> Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: vthumbe1503 <[email protected]> Co-authored-by: Varun Thumbe <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Vladimir Cherepanov <[email protected]>
…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 <[email protected]> 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 <[email protected]> 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 <[email protected]> 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 <[email protected]> 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 <[email protected]> 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 <[email protected]> Made-with: Cursor * tests skip Signed-off-by: Pawel Gadzinski <[email protected]> * tests skip Signed-off-by: Pawel Gadzinski <[email protected]> * Warn on deprecated kwargs in Float8CurrentScalingQuantizer Signed-off-by: Pawel Gadzinski <[email protected]> * 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 <[email protected]> --------- Signed-off-by: Pawel Gadzinski <[email protected]> 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 <[email protected]>
… (#2728) * fix: scope get_full_cu_seqlens cache key by device and inference mode Signed-off-by: Dongmin Ra <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Dongmin Ra <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Charlene Yang <[email protected]>
Fix FA4 selection when FA3 is unavailable. Signed-off-by: Björn Buschkämper <[email protected]>
* 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 <[email protected]> Signed-off-by: dcampora <[email protected]> Signed-off-by: Daniel Campora <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <[email protected]> * 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 <[email protected]> Co-Authored-By: Claude Sonnet 4.6 <[email protected]> Signed-off-by: Daniel Campora <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <[email protected]> * 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 <[email protected]> Signed-off-by: dcampora <[email protected]> Signed-off-by: Daniel Campora <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <[email protected]> * 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 <[email protected]> Signed-off-by: Daniel Campora <[email protected]> * Add assertion for 2D input in topk Co-Authored-By: Claude Sonnet 4.6 <[email protected]> Signed-off-by: Daniel Campora <[email protected]> --------- Signed-off-by: dcampora <[email protected]> Signed-off-by: Daniel Campora <[email protected]> Signed-off-by: Diego Campora <[email protected]> Co-authored-by: Claude Sonnet 4.6 <[email protected]> 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 <[email protected]>
* remove ctype to eliminate memory usage from the cudnn kernel Signed-off-by: Varun Thumbe <[email protected]> * Remove c_dtype from fusible ops test Signed-off-by: Kirthi Shankar Sivamani <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]> Signed-off-by: Kirthi Shankar Sivamani <[email protected]> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
Signed-off-by: Muu <[email protected]>
* [Common, PyTorch] Add triton mHC kernels & pytorch operators Signed-off-by: Kaining Zhong <[email protected]> * fix Signed-off-by: Kaining Zhong <[email protected]> * nit Signed-off-by: Kaining Zhong <[email protected]> * make linter happy Signed-off-by: Kaining Zhong <[email protected]> * nit Signed-off-by: Kaining Zhong <[email protected]> * ah OK Signed-off-by: Kaining Zhong <[email protected]> * new configs to improve perf Signed-off-by: Kaining Zhong <[email protected]> * add APIs to docs Signed-off-by: Kaining Zhong <[email protected]> * fix typos, check deterministic, refactor Signed-off-by: Kaining Zhong <[email protected]> * fix Signed-off-by: Kaining Zhong <[email protected]> * [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 <[email protected]> * add docstring Signed-off-by: Kaining Zhong <[email protected]> * fix api doc Signed-off-by: Kaining Zhong <[email protected]> * whoops Signed-off-by: Kaining Zhong <[email protected]> * grad_x doesn't have to zero Signed-off-by: Kaining Zhong <[email protected]> * nit Signed-off-by: Kaining Zhong <[email protected]> * nit Signed-off-by: Kaining Zhong <[email protected]> * force pytorch to not use bf16 for reduction Signed-off-by: Kaining Zhong <[email protected]> * use TE's general_gemm instead Signed-off-by: Kaining Zhong <[email protected]> * [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 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kaining Zhong <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…True (#2936) * fix Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup Signed-off-by: Varun Thumbe <[email protected]> * [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 <[email protected]> --------- Signed-off-by: Varun Thumbe <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root <[email protected]>
… (#2924) * Fix contiguous path for k=2880 Signed-off-by: Kirthi Shankar Sivamani <[email protected]> * format Signed-off-by: Kirthi Shankar Sivamani <[email protected]> * Review suggestion from @Oleg-Goncharov Signed-off-by: Kirthi Shankar Sivamani <[email protected]> * Add test for swizzle + padding fusion Signed-off-by: Kirthi Shankar Sivamani <[email protected]> * Address review comments Signed-off-by: Kirthi Shankar Sivamani <[email protected]> --------- Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
…942) * accumulate bias in fp32 instead of bf16 in ref impl dbias to avoid accumulated numerical error Signed-off-by: tdophung <[email protected]>
* loosen up thresholds. Also only check min loss of last 10% of steps to avoid failing by noise near convergence Signed-off-by: tdophung <[email protected]> * add deterministic flag for mnist run Signed-off-by: tdophung <[email protected]> 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 <[email protected]> * Added confirmation with uniformity in one of the dimensions Signed-off-by: Abhishek <[email protected]> * [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 <[email protected]> * [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 <[email protected]> * Updated the code with newer changes in main Signed-off-by: Abhishek <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Abhishek <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 <[email protected]>
…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 <[email protected]> * Tweak test to catch bug when alternating train and infer steps Signed-off-by: Tim Moon <[email protected]> * Avoid removing usages from quantized weights in grouped linear op Signed-off-by: Tim Moon <[email protected]> * [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 <[email protected]> * Blindly preserve quantizer usages in quantized weight params. Signed-off-by: Tim Moon <[email protected]> --------- Signed-off-by: Tim Moon <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
… (#2943) * Patch for NVFP4 test suite Signed-off-by: Oleg Goncharov <[email protected]> * C++ tests fix Signed-off-by: Oleg Goncharov <[email protected]> * Cleanup Signed-off-by: Oleg Goncharov <[email protected]> * [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 <[email protected]> * Set the golden value for amax in tests Signed-off-by: Oleg Goncharov <[email protected]> * Fixed memory leakage Signed-off-by: Oleg Goncharov <[email protected]> --------- Signed-off-by: Oleg Goncharov <[email protected]> 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 <[email protected]> * Fix tests Signed-off-by: Kirthi Shankar Sivamani <[email protected]> --------- Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
…n expert (#2947) Add workaround for cuteDSL stride requirement for zero token expert Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
Remove unused NVFP4 quantize kernel Signed-off-by: Tim Moon <[email protected]> Co-authored-by: Oleg Goncharov <[email protected]>
Follow-up to bffa8f0, trimmed to the necessary changes: - quantization_ref_nvfp4.py: keep only the be3feea revert (global scale); restore the incidental is_fp8_fnuz cleanup (no-op on gfx950, unreached on gfx942) to keep the diff minimal. - test_dequantize_nvfp4.cu: skip on gfx942 (te_fp8_fnuz) — the dequant kernel now uses OCP 448 while its reference is fnuz-240, so it fails there. - test_nvfp4_rht_quantize_exact.py: gate the ROCm-only test_hadamard_transform_amax on recipe_available (was IS_HIP_EXTENSION-only; uses the reverted reference). Drops the unnecessary cast/grouped-gemm cpp skips (those tests pass on gfx942). Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
68bc6e0 added the te_fp8_fnuz() skip to the CUDA-branch dequant bodies (#else of __HIP_PLATFORM_AMD__), which don't run on ROCm — so gfx942 still failed — and would break the CUDA build (te_fp8_fnuz is ROCm-only). Move the skip to the ROCm-branch DequantizeNVFP4 test body (the one that actually runs on gfx942) and drop it from the CUDA bodies. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
|
@ipanfilo @wangye805 related to your comment on I cannot comment there but ill add here my reply: I replaced the TODO with an explanatory comment in the |
ipanfilo
left a comment
There was a problem hiding this comment.
Number of files have AMD copyright header incorrect format
| import transformer_engine_torch as tex | ||
|
|
||
| from transformer_engine.common.recipe import Recipe | ||
| from transformer_engine.pytorch.tensor._quantization_helpers import ( |
There was a problem hiding this comment.
So on ROCm _quantized_helpers import quntized_tensor? Why?
|
@ipanfilo In response to your question about The cycle is ROCm-specific because |
|
@ipanfilo I fixed all remaining AMD copyright and added any missing ones. |
new_empty([]) crashed staging a rank-0 NVFP4 placeholder; shape it like the source instead. The per-tensor amax rides fsdp_pre_all_gather metadata, which async DCP bf16-casts, so async cannot be bit-exact -- use the documented loose tolerance (as with Float8CurrentScaling). Sync stays bit-exact. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
TE caches the FP8 weight columnwise transpose during forward instead of recreating it in backward, accumulating ~0.68 MiB/layer under FSDP2. The modules already skip caching when use_fsdp2 is set, which the test never did; enable use_fsdp2 on submodules in _shard_model after sharding. Also trims the verbose new_empty comment added in the prior NVFP4 async-DCP fix. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
# Conflicts: # transformer_engine/common/fused_attn_rocm/fused_attn.cpp
|
| __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const float global_amax) { | ||
| using namespace detail; | ||
| constexpr float fp8_max = TypeExtrema<fp8e4m3>::max; // 448.0f; | ||
| static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); | ||
| constexpr float fp8_max = static_cast<float>(E4M3_MAX); |
There was a problem hiding this comment.
Are any of those changes ROCm specific or all come from upstream?
There was a problem hiding this comment.
I reverted the copyright since this file is all from upstream.
|
@ipanfilo During my check of copyright I found a couple of files not touched by this IFU but which had AMD content but the copyright was missing: I checked again and now everything else should be clean with the copyrights. |
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: