Upgrade/v2.17 gated 20260810 - #101
Open
zhaoyinglia wants to merge 73 commits into
Open
Conversation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
…A#2274) * Fix imports in test for deprecated jax.experimental.pjit 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 * Fix: Pass NamedSharding instead of PartitionSpec to compare_ops() so that when the in and out sharding is used to create a jitted function, it has the mesh info Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.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-eos01.eos.clusters.nvidia.com>
* Support wheel build for cuda 13 Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fixes Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fixes for cu13 runtime, format Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Add documentation Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Better error handling Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * fix Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * fix jax sdist Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Modify function names Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…tization (NVIDIA#2270) * [JAX] Support recipe flags for disabling SR, RHT, and 2D quantization Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix issue with SR state being erased due to pytree handling of NVFP4Quantizer Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Add test for SR state preservation across VJP boundaries Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix sharding of SR rng state Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * update tolerances slightly now that SR is enabled Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Use hashlib for deterministic hashes across runs for SR Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * rename uses_rht on scaled tensors to has_applied_rht Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * add assert Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Move decision of whether to use RHT into helper.py and add dedicated RHT tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * fix use_rht attr usage Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * fix pure-jax rht usage criteria Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Adjust tolerances after rebase Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> --------- Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Include TE core headers in build Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Added sm_120f to the build Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Change the arch specific handling Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Support for CUDA<12.9 Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Moved through the rest of the files Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Common cases Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Remove pure 100 from the list Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * CMake changes, (not yet working) Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Do not pass the arch-specific thing from build_tools Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Moved some of the files to arch-specific compilation Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix and also changing the order of compilation to hopefully get the compilation time lower Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix for the files overwriting custom compile properties Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Actually make this whole thing work Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add space to the error message Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> * Apply suggestions from code review Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> * Fixes from review Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Changing the naming to be more intuitive Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add missing cassert include for device-side asserts Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
* add max_score for fused/unfused F16 non-CP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * calculate max per head instead of max over all heads Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fused attn max_score shape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert FE to github 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 to 1.15.0-rc Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * 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 * reduce ew kernels; fix causal masks; add more 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 * minor fix to 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 * remove logic for flash-attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: add CP support for p2p/a2a/all_gather 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 improvements of implementation/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 * WIP: add thd support Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add thd to UnfusedDPA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more fixes for lint 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 to FE 1.15 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unneeded changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable unfused for thd + pad_between_seqs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable thd for unfused until bug is fixed 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 all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * rename max_score to max_logit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable fused attn + thd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- 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>
…DIA#2288) * Fix CI failures due to deterministic attention Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * some more cleanup Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix debug test Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…2315) * Fix: Skip determinism tests for bprop for all sm >=100 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add username to TODO Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Assert in fused attn bwd pass for sm100+ 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 --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Fix attention backend and tests for sm120 Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Disable MLA only for backward Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…determinism on Blackwell (NVIDIA#2316) * Bump the min version expected to supported FP8 cs det on Blackwell Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Disable fused attn for cudnn < 9.14 for FP8 CS. Disable fused attn for cudnn < 9.18 for FP8 deterministic CS 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 --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…NVIDIA#2322) Ensure JAX reference impl uses an accurate backend Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
# Description Add the FlagOS multi-chip backend for TransformerEngine Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Co-authored-by: zhaoyinglia <ylzhao@baai.ac.cn>
# Description Fix import bugs Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
# Description Please include a brief summary of the changes, relevant motivation and context. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
…VIDIA#4) # TransformerEngine-FL Plugin System ## Overview This PR implements a comprehensive multi-backend plugin system for TransformerEngine-FL, enabling support for multiple hardware vendors (NVIDIA, AMD, Hygon, etc.) while maintaining full API compatibility with the original `transformer_engine_torch`. **Core Philosophy**: A plugin-based backend system that allows hardware vendors to easily implement their own operator optimizations while preserving complete compatibility with the original TransformerEngine API. ## Key Features ### Full API Compatibility - Drop-in replacement for `transformer_engine_torch` - Switch backends via environment variables - Zero changes required to existing user code ### Multi-Backend Support | Backend | Description | Implementation | |---------|-------------|----------------| | **FlagOS (default)** | Triton-based cross-platform implementation | `backends/flagos/` | | **CUDA (vendor)** | Wraps original TransformerEngine C++ extensions | `backends/vendor/cuda/` | | **Reference** | Pure PyTorch fallback implementation | `backends/reference/` | ### Three-Tier Backend Selection ``` ┌─────────────────────────────────────────────────────────┐ │ 1. TE_FL_PER_OP (Per-operator override) [Highest] │ │ Example: TE_FL_PER_OP="rmsnorm_fwd=vendor:cuda" │ ├─────────────────────────────────────────────────────────┤ │ 2. TE_FL_PREFER (Global preference) │ │ Values: flagos / vendor / reference │ ├─────────────────────────────────────────────────────────┤ │ 3. Backend Priority (Intrinsic) [Lowest] │ │ Each implementation has a priority value │ └─────────────────────────────────────────────────────────┘ ``` ## Architecture ### Directory Structure ``` transformer_engine/plugin/core/ ├── __init__.py # Public API exports ├── types.py # Core types: BackendImplKind, OpImpl ├── registry.py # OpRegistry: stores all implementations ├── manager.py # OpManager: selects and calls implementations ├── policy.py # SelectionPolicy: backend selection rules ├── discovery.py # Plugin auto-discovery (entry_points, env) ├── builtin_ops.py # Registers all built-in backends ├── ops.py # TEFLModule: transformer_engine_torch compatible API ├── logger_manager.py # Logging utilities ├── _module_setup.py # Module aliasing setup ├── _build_config.py # Build-time configuration │ └── backends/ ├── flagos/ # FlagOS backend (Triton-based) │ ├── flagos.py # FlagOSBackend class │ ├── register_ops.py # Operator registration │ └── impl/ # Operator implementations │ ├── rmsnorm.py │ ├── gemm.py │ └── ... │ ├── vendor/ # Vendor backends │ └── cuda/ # NVIDIA CUDA backend │ ├── cuda.py # CUDABackend class │ └── register_ops.py │ └── reference/ # Reference backend (PyTorch) ├── reference.py # ReferenceBackend class ├── register_ops.py └── impl/ # Pure PyTorch implementations ``` ### Core Components | File | Description | |------|-------------| | `types.py` | Defines `BackendImplKind` (DEFAULT/VENDOR/REFERENCE) and `OpImpl` dataclass | | `registry.py` | `OpRegistry` - Central storage for all operator implementations | | `manager.py` | `OpManager` - Handles implementation selection, fallback, and execution | | `policy.py` | `SelectionPolicy` - Configurable rules for backend selection | | `discovery.py` | Auto-discovers plugins via `entry_points` or `TE_FL_PLUGIN_MODULES` | | `ops.py` | `TEFLModule` - Provides `transformer_engine_torch` compatible interface | ## Installation ### Build with CUDA support ```bash pip install --no-build-isolation -e . ``` ### Build without CUDA (FlagOS only) ```bash TE_FL_SKIP_CUDA=1 pip install --no-build-isolation -e . ``` ## Environment Variables ### Backend Selection | Variable | Description | Values | Default | |----------|-------------|--------|---------| | `TE_FL_PREFER` | Preferred backend type | `flagos` / `vendor` / `reference` | `flagos` | | `TE_FL_PREFER_VENDOR` | Prefer vendor (legacy) | `1` / `0` | `0` | | `TE_FL_STRICT` | Strict mode (no fallback) | `1` / `0` | `0` | ### Vendor Filtering | Variable | Description | Example | |----------|-------------|---------| | `TE_FL_ALLOW_VENDORS` | Allowed vendors (whitelist) | `nvidia,amd` | | `TE_FL_DENY_VENDORS` | Denied vendors (blacklist) | `vendor_a` | ### Per-Operator Configuration | Variable | Description | Example | |----------|-------------|---------| | `TE_FL_PER_OP` | Per-operator backend ordering | `rmsnorm_fwd=vendor:cuda\|default` | ### Plugin Discovery | Variable | Description | Example | |----------|-------------|---------| | `TE_FL_PLUGIN_MODULES` | Plugin modules to load | `my_plugin,another_plugin` | ### Build Configuration | Variable | Description | Values | Default | |----------|-------------|--------|---------| | `TE_FL_SKIP_CUDA` | Skip CUDA backend | `1` / `0` | `0` | | `CUDA_HOME` | CUDA installation path | `/usr/local/cuda` | Auto-detected | ### Logging | Variable | Description | Values | Default | |----------|-------------|--------|---------| | `TEFL_LOG_LEVEL` | Log level | `DEBUG` / `INFO` / `WARNING` / `ERROR` | `INFO` | ## Usage Examples ### Basic Usage (No Code Changes Required) ```python # Existing code works as-is import transformer_engine.pytorch as te # or import transformer_engine_torch as te ``` ### Register Custom Backend (In-tree) ```python from transformer_engine.plugin.core import ( OpRegistry, OpManager, OpImpl, BackendImplKind ) # 1. Define implementation def my_rmsnorm(input, weight, eps=1e-5, **kwargs): variance = input.pow(2).mean(-1, keepdim=True) return input * torch.rsqrt(variance + eps) * weight, torch.rsqrt(variance + eps) # 2. Register registry = OpRegistry() registry.register_impl(OpImpl( op_name="rmsnorm_fwd", impl_id="vendor.mybackend", kind=BackendImplKind.VENDOR, vendor="mybackend", fn=my_rmsnorm, priority=200, )) # 3. Call manager = OpManager(registry) output, rsigma = manager.call("rmsnorm_fwd", input, weight) ``` ### Register Custom Backend (Out-of-tree Plugin) Create a plugin package with `register(registry)` function: ```python # my_vendor_plugin/__init__.py from transformer_engine.plugin.core import OpImpl, BackendImplKind def my_rmsnorm(input, weight, eps=1e-5, **kwargs): # Your implementation ... def register(registry): """Called automatically by TE-FL""" registry.register_impl(OpImpl( op_name="rmsnorm_fwd", impl_id="vendor.myvendor", kind=BackendImplKind.VENDOR, vendor="myvendor", fn=my_rmsnorm, priority=200, )) ``` Load via environment variable: ```bash export TE_FL_PLUGIN_MODULES=my_vendor_plugin python your_script.py ``` ## Runtime Logs When running, you'll see logs indicating which backend is used: ``` [TE-FL manager.py:133 INFO] Registered impl_ids: ['default.flagos', 'reference.torch', 'vendor.cuda'] [TE-FL manager.py:390 INFO] Op 'rmsnorm_fwd' using 'default.flagos' (kind=default, vendor=None) [TE-FL manager.py:395 INFO] Op 'rmsnorm_fwd' switched from 'default.flagos' to 'vendor.cuda' (kind=vendor, vendor=CUDA) ``` ## Examples See `transformer_engine/plugins/examples/` for complete working examples: - `example_intree.py` - In-tree backend registration - `example_outtree.py` - Out-of-tree plugin registration Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Co-authored-by: panpy <panpy@sugon.com>
# Description - Add missing __init__.py files to transformer_engine/plugin/core/backends/flagos/attention/ directory tree to fix import errors when accessing these modules as Python packages - Add comprehensive test suite (test_policy.py) covering the TE-FL scheduling policy system including: SelectionPolicy creation and configuration Environment variable parsing (TE_FL_PREFER, TE_FL_STRICT, etc.) Policy context managers Vendor filtering (allow/deny) Thread safety validation Minor code style improvements Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
…n fallback (NVIDIA#14) ## Summary This PR contains two major improvements: 1. **Register `get_attention_backend` function for all backends** (CUDA, FlagOS, Reference) - Added `get_attention_backend` implementation to all backend types - Ensures consistent attention backend selection across different hardware platforms 2. **Fix FlashAttention fallback mechanism** - Removed redundant `_called_impls` dictionary, replaced with simpler `_last_impl_id` class variable - Removed unused `_log_lock` threading lock - Simplified implementation tracking and logging logic - Reduced code complexity and memory overhead while maintaining full functionality ## Changes - Updated `FlashAttentionBase` class in `ops.py` to remove redundant implementation tracking - Added `get_attention_backend` registration to CUDA, FlagOS, and Reference backends - Fixed fallback logic in attention backend selection ## Test Plan - [x] Code builds successfully - [x] Existing tests pass - [x] Manual testing with different backend configurations ## Related Issues Fixes issues with FlashAttention fallback and improves backend consistency.
# Description fix nv shared lib bug [CUDA] Import failed: No module named 'transformer_engine_torch_nv' Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
# Description This pr add hygon backend for calling basic ops on hygon dcu. ## Type of change - [x] New feature (non-breaking change which adds functionality) ## Changes Please list the changes introduced in this PR: - Add a new `hygon` folder in `vendor` contains `__init__.py`, `hygon.py`, `register_ops.py` - Register hygon ops in `builtin_ops.py` # Requirements In order to use hygon backend, the following, the following requirements need to be met - The python package `transformer_engine_fl_hygon` needs to be installed # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Signed-off-by: wenjh <wenjh@sugon.com>
…DIA#18) Add a flag that permanently enables flag_gems with a single switch, eliminating the need to call flag_gems.use_gems for every single operator. This removes significant registration overhead and improves end-to-end throughput. - When the flag is set, every operator’s implementation is forced to use flag_os/vendor; the default PyTorch reference backend is unavailable. - When the flag is not set, operators can freely switch among flag_os, vendor, and torch backends.
Unify the usage of the gems context - only enter or exit the context when switching between the flagos backend and the torch backend (or vice versa). - avoids the overhead of repeated enter/exit calls across multiple OPs.
## Summary - Support combined qkv_layout formats like `sbhd_sbhd_sbhd` by extracting the first part for layout conversion - Distinguish between standard 4D tensor format (sbhd/bshd) and true packed format (thd). For 4D tensors, directly convert layout like flagos backend does, instead of incorrectly trying to unpack ## Problem When using torch SDPA backend with `batch_size > 1`, the following error occurs: ``` ValueError: Unexpected 4D tensor shape torch.Size([4096, 4, 16, 128]). Expected [total_tokens, 1, num_heads, head_dim] ``` The original code incorrectly tried to unpack 4D tensors when `cu_seqlens` was provided, but 4D tensors in `sbhd`/`bshd` format should be handled with simple layout conversion (like flagos backend does). ## Test plan - [x] Tested with batch_size=4, verified no ValueError - [x] Results match flagos backend output
- Remove the flag_gems.use_gems() context to avoid context-switching overhead - Call flag_gems.xxx directly wherever possible.
# Description Add the new vendor backend METAX ## Type of change - [ ] New feature (non-breaking change which adds functionality) ## Changes Please list the changes introduced in this PR: - Add metax ops register - Add metax backend implementation - Register metax ops in builtin_ops.py ## Requirements - The module transformer_engine_torch_metax is needed, to use this module, need to install package transformer_engine_metax # Checklist: - [x] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes
…VIDIA#23) ## Summary - flagos: Add multi_tensor_adam_param_remainder implementation - reference: Add multi_tensor_adam_param_remainder implementation - reference: Add context parallel support for Flash Attention - manager: Add cache mechanism with _impl_cache and _impl_cache_meta for conditional op selection ## Changes ### flagos backend - Implemented multi_tensor_adam_param_remainder operation for handling parameter remainders in multi-tensor Adam optimizer ### reference backend - Implemented multi_tensor_adam_param_remainder operation - Added context parallel support for Flash Attention implementation ### Core manager - Added cache mechanism using _impl_cache and _impl_cache_meta - Improved op selection with conditional caching based on policy fingerprint and epoch --------- Signed-off-by: wenone766 <wenone766@wenobug.com> Co-authored-by: wenone766 <wenone766@wenobug.com>
- Fix enum mismatch, between ```transformer_engine/plugin/core/ops.py``` and ```transformer_engine/common/include/transformer_engine/xxx.h```
# Description add Vendor KUNLUNXIN
Fix kunlunxin register errors
# Description
Add the new vendor backend ENFLAME
## Type of change
- [ √ ] New feature (non-breaking change which adds functionality)
## Changes
Please list the changes introduced in this PR:
- Add enflame ops register
- Add enflame backend implementation
- Register enflame ops in builtin_ops.py
## Requirements
- The module migraiton is needed, to use this module, need to install
package migration whl
# Checklist:
- [x] I have read and followed the [contributing
guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst)
- [x] The functionality is complete
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
Co-authored-by: xijun.gong <xijun.gong@enflame-tech.com>
# Description Add some kunlunxin ops bind code ## Type of change - [1] New feature (non-breaking change which adds functionality) ## Changes Add kunlunxin backend bind support. Add kunlunxin ops bind and register. # Checklist: - [1] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [1] The functionality is complete - [1] I have commented my code, particularly in hard-to-understand areas - [1] I have made corresponding changes to the documentation - [1] My changes generate no new warnings - [1] I have added tests that prove my fix is effective or that my feature works - [1] New and existing unit tests pass locally with my changes
**Merge upstream release_v2.14 f031cf8 into main** Integrates NVIDIA TransformerEngine upstream release v2.14 (304 commits, v2.9.0 → v2.14.0) into the TransformerEngine-FL fork via tree replacement strategy, preserving the custom plugin system while incorporating upstream enhancements. **Upstream Enhancements** Quantization & Precision - MXFP8 grouped GEMM with persistent quantization kernels and tensor-scaled FP8 support - NVFP4 grouped quantization with Hadamard transform for MoE workloads - QuantizedTensor support in FusedAdam optimizer for MXFP8/Float8 block scaling Architecture Support - Blackwell (sm120) fused attention support with cuDNN 9.18.1+ - Deterministic training on Blackwell with cuDNN ≥9.18.1 - Grouped GEMM cuBLAS bindings with bias support and tensor swizzling Distributed Training - FSDP2 support with DTensor-aware optimizer states and allgather optimizations - Collective GEMM with FP8/MXFP8 for JAX - GroupedTensor torch ops for DDP and distributed optimizer Operators - Fused RMSNorm dLN with add-through via cuDNN - MoE grouped MLP ops with split dBias and router kernel JAX bindings - Configurable philox rounds for stochastic rounding **FlagOS Features** Plugin System Preservation - Synced plugin OP API signatures with upstream csrc changes (fused_attn_fwd/bwd parameters, attention backend dispatch) - Patched new upstream CUDA hardcoding to te_device_type() for multi-backend compatibility - Fixed stale references to renamed upstream symbols (e.g., CPUOffloadEnabled → is_cpu_offload_enabled()) Verification - Build & import validation passed - Unit & integration tests passed - FlagScale end to end training test, summary as follows: **Qwen3-32B, 16 layers, 20 iters, 1node x 8 gpus** | Config | Status | Avg Throughput (tokens/s/gpu) | Note | |------|------|------|------| | vendor-flash | PASS | 124.33 | | | vendor-fused | PASS | 121.95 | | | vendor-unfused | PASS | 108.35 | | | flagos-flash | PASS | 94.20 | | | flagos-fused | FAIL | — | No fused attention backend supports for flagos backend | | flagos-unfused | PASS | 65.01 | | | reference-flash | PASS | 93.14 | | | reference-fused | FAIL | — | No fused attention backend support for reference backend | | reference-unfused | PASS | 66.89 | | **DeepSeek-V3 16BA3B, 18 layers with 1 mtp layer, 20 iters, 1node x 8gpus, there is no flash-attn or fused-attn support for multi-latent attention** | Config | Status | Avg Throughput (tokens/s/gpu) | Note | |------|------|------|------| | vendor-unfused | PASS | 47.00 | | | flagos-unfused | PASS | 18.57 | | | reference-unfused | PASS | 20.68 | | --------- Signed-off-by: Jack <lityangweiguang@163.com> Signed-off-by: oliver könig <okoenig@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: janbernloehr <jan@bernloehrs.de> Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: tdophung <tdophung@nvidia.com> Signed-off-by: Shoval Atias <satias@satias-mlt.client.nvidia.com> Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Signed-off-by: Pingtian Li <pingtianl@nvidia.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Signed-off-by: Evgeny <etsykunov@nvidia.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> Signed-off-by: kunlunl <kunlunl@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com> Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-preos01.a51.clusters.nvidia.com> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Robin Zhang <robinz@nvidia.com> Signed-off-by: ykarnati <ykarnati@nvidia.com> Signed-off-by: Keith Wyss <kwyss@nvidia.com> Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> Signed-off-by: Jinhang Choi <jinhangc@nvidia.com> Signed-off-by: LucienXian <fl.xian@foxmail.com> Signed-off-by: xiaoxi-wangfj <690912414@qq.com> Signed-off-by: fuyue.lj <fuyue.lj@antgroup.com> Signed-off-by: Peter St. John <pstjohn@nvidia.com> Signed-off-by: Victor Oliveira <victor.oliveira@getcruise.com> Signed-off-by: hongbinl <hongbinl@nvidia.com> Signed-off-by: Hongbin Liu <lhb8125@users.noreply.github.com> Signed-off-by: Hongbin Liu <hongbinl@nvidia.com> Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> Signed-off-by: Chen Cui <chcui@nvidia.com> Signed-off-by: DoubleCheeseCheetos <hanhdp99@gmail.com> Signed-off-by: tdophung <hanhdp99@gmail.com> Signed-off-by: Piotr Gadzinski <pgadzinski@nvidia.com> Signed-off-by: Vadim Markovtsev <vadim@poolside.ai> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> Signed-off-by: Lifu Zhang <lifuz@login-lyris02.lyris.clusters.nvidia.com> Signed-off-by: Kim, Jin <jinn.kim@sk.com> Signed-off-by: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com> Signed-off-by: JAX Toolbox <jax@nvidia.com> Signed-off-by: Hemil Desai <hemild@nvidia.com> Signed-off-by: Xin Yao <xiny@nvidia.com> Signed-off-by: Nicolas Castet <ncastet@nvidia.com> Signed-off-by: Alp Dener <adener@nvidia.com> Signed-off-by: Gao <gdeng@nvidia.com> Signed-off-by: Xin Yao <yaox12@outlook.com> Signed-off-by: tongliu <tongliu@nvidia.com> Signed-off-by: root <root@eos0046.eos.clusters.nvidia.com> Signed-off-by: qiyuw <qiyuw@nvidia.com> Signed-off-by: aagallo <aagallo@amazon.com> Signed-off-by: Andrea Gallo <aagallo@amazon.com> Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com> Signed-off-by: Chaoyang Mei <1192554423@qq.com> Signed-off-by: meichaoyang001 <meichaoyang001@ke.com> Signed-off-by: Sung Hyun Cho <hope5487@gmail.com> Signed-off-by: Bias92 <pewpewplay315@gmail.com> Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com> Signed-off-by: Zhiyi Su <dantesuu@gmail.com> Signed-off-by: ZhiyiDanielSu <35579247+zobeideThePlayer@users.noreply.github.com> Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com> Signed-off-by: Jonathan Mitchell <jomitchell@umb-b300-dp-147.ipp4a1.colossus.nvidia.com> Signed-off-by: Jonathan Mitchell <jomitchell@ipp1-1429.ipp1a1.colossus.nvidia.com> Signed-off-by: Peter St. John <peterc.stjohn@gmail.com> Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt> Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Signed-off-by: Cory Ye <cye@nvidia.com> Signed-off-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: Jack <lityangweiguang@163.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Jan Bernlöhr <jan@bernloehrs.de> Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: satias10 <satias@nvidia.com> Co-authored-by: Shoval Atias <satias@satias-mlt.client.nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Pingtian Li <158665726+Wohox@users.noreply.github.com> Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: Evgeny Tsykunov <etsykunov@nvidia.com> Co-authored-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: Kunlun Li <94586211+kunlunl@users.noreply.github.com> Co-authored-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com> Co-authored-by: Kshitij Janardan Lakhani <klakhani@login-preos01.a51.clusters.nvidia.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: Ming Huang <mingh@nvidia.com> Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: Robin Zhang <robinz@nvidia.com> Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Co-authored-by: kwyss-nvidia <kwyss@nvidia.com> Co-authored-by: vcherepanov-nv <vcherepanov@nvidia.com> Co-authored-by: Jinhang Choi <cepiross@gmail.com> Co-authored-by: LucienXian <fl.xian@foxmail.com> Co-authored-by: xiaoxi-wangfj <690912414@qq.com> Co-authored-by: 刘俊 <fuyue.lj@antgroup.com> Co-authored-by: Peter St. John <pstjohn@nvidia.com> Co-authored-by: Victor Oliveira <victormatheus@gmail.com> Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com> Co-authored-by: Santosh Bhavani <santosh.bhavani@live.com> Co-authored-by: Jacket <44538064+kainzhong@users.noreply.github.com> Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Co-authored-by: Chen Cui <chcui@nvidia.com> Co-authored-by: DoubleCheeseCheetos <hanhdp99@gmail.com> Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com> Co-authored-by: Vadim Markovtsev <gmarkhor@gmail.com> Co-authored-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> Co-authored-by: Lifu Zhang <tomzhanglf@gmail.com> Co-authored-by: Lifu Zhang <lifuz@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Zhongbo Zhu <zhongboz@nvidia.com> Co-authored-by: Kim, Jin (Jay@SKT) <singleheart@gmail.com> Co-authored-by: Harikrishna KP <harikp2002@gmail.com> Co-authored-by: JAX Toolbox <jax@nvidia.com> Co-authored-by: Hemil Desai <hemil.desai10@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: Nicolas Castet <26874160+nvcastet@users.noreply.github.com> Co-authored-by: Alp Dener <adener@nvidia.com> Co-authored-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: Gao <gdeng@nvidia.com> Co-authored-by: Tong Liu <tongliu@nvidia.com> Co-authored-by: root <root@eos0046.eos.clusters.nvidia.com> Co-authored-by: root <root@eos0022.eos.clusters.nvidia.com> Co-authored-by: Qiyu Wan <39144338+WanZzzzzz@users.noreply.github.com> Co-authored-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: aagallo <aagallo@amazon.com> Co-authored-by: aagallo <aagallo@amzon.com> Co-authored-by: Fabian Joswig <fjosw@users.noreply.github.com> Co-authored-by: Chaoyang Mei <1192554423@qq.com> Co-authored-by: Sung Hyun Cho <hope5487@gmail.com> Co-authored-by: 노란토끼 <83907395+Bias92@users.noreply.github.com> Co-authored-by: vasunvidia <108759426+vasunvidia@users.noreply.github.com> Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Zhiyi Su <dantesuu@gmail.com> Co-authored-by: ZhiyiDanielSu <35579247+zobeideThePlayer@users.noreply.github.com> Co-authored-by: jomitchellnv <148147880+jomitchellnv@users.noreply.github.com> Co-authored-by: Jonathan Mitchell <jomitchell@ipp1-1334.ipp1a1.colossus.nvidia.com> Co-authored-by: Jonathan Mitchell <jomitchell@umb-b300-dp-147.ipp4a1.colossus.nvidia.com> Co-authored-by: Jonathan Mitchell <jomitchell@ipp1-1429.ipp1a1.colossus.nvidia.com> Co-authored-by: Jeremy Berchtold <jberchtold@nvidia.com> Co-authored-by: Carlos Gomes <carlosmiguel.gomes@live.com.pt> Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: lixianduo <lixianduo@flagos.ai> Co-authored-by: BrianPei <kaworu228@gmail.com> Co-authored-by: qqjxzxq <1376782660@qq.com> Co-authored-by: HermiaHuan <3081497279@qq.com>
Fix op register errors when skip cuda, add some optimizer ops for reference backend
Support bias for generic_gemm
add flagos layernorm
# Description Added the binding and invocation for the `multi_tensor_compute_scale_inv_e8m0` operator for the Kunlunxin vendor, and corrected a syntax error in the invocation of the `multi_tensor_compute_scale_and_scale_inv` operator. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [1] Bug fix (non-breaking change which fixes an issue) - [1] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Added code for the binding and invocation of the `multi_tensor_compute_scale_inv_e8m0` operator. Modified the code calling `multi_tensor_compute_scale_and_scale_inv`. # Checklist: - [ 1] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ 1] The functionality is complete - [ 1] I have commented my code, particularly in hard-to-understand areas - [ 1] I have made corresponding changes to the documentation - [ 1] My changes generate no new warnings - [ 1] I have added tests that prove my fix is effective or that my feature works - [ 1] New and existing unit tests pass locally with my changes
# Description Please include a brief summary of the changes, relevant motivation and context. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
## Summary
Replace the grouped GEMM plugin test import from
`transformer_engine_torch_nv` to `transformer_engine_torch`.
## Why
The test should load the standard Transformer Engine torch extension
module. Importing `transformer_engine_torch_nv` can fail because that is
not the expected module name in this project.
## Impact
This fixes startup/import failure for
`transformer_engine/plugin/tests/test_te_general_grouped.py` before the
grouped GEMM checks run.
## Validation
- `git diff --check origin/main...HEAD`
- `python3 -c 'import ast, pathlib; path =
pathlib.Path("transformer_engine/plugin/tests/test_te_general_grouped.py");
ast.parse(path.read_text())'`
Full test execution was not run locally because it requires the project
runtime/CUDA environment.
NVIDIA#82) # Description When loading Hygon libraries, the code previously assumed that `hygon_spec.origin` is always available and directly accessed it to determine the package path. However, for namespace packages or certain installation layouts, `importlib.util.find_spec()` may return a spec whose `origin` is `None`, causing an exception when accessing `Path(hygon_spec.origin)`. This PR adds a fallback mechanism: - Use `hygon_spec.origin` when available. - Fall back to `hygon_spec.submodule_search_locations` when `origin` is `None`. - Return gracefully with an error message if neither source can provide a valid package path. This improves compatibility with different Python package layouts and prevents startup failures when loading Hygon-related libraries. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Add a null check for `hygon_spec.origin` - Fall back to `hygon_spec.submodule_search_locations[0]` when `origin` is unavailable - Add explicit error handling when neither `origin` nor `submodule_search_locations` can determine the package path - Prevent crashes caused by `Path(None)` during Hygon library loading # Checklist: - [x] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes --------- Co-authored-by: wangyl <wangyl16@sugon.com> Co-authored-by: wangyl166 <601199939@qq.com>
…odule) on FlagOS (NVIDIA#79) # Description This PR optimizes the training entry of FlagScale, enabling native NPU training capability without introducing MegatronAdaptor dependencies. 1、Retain complete distributed training ability on NPU 2、No breaking changes to existing FlagScale training workflows Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
# Description This PR focuses on improving unit test coverage for the plugin core framework and expanding CI test coverage by re-enabling previously excluded test suites. A coverage analysis was performed on the current CI workflow to identify uncovered code paths and low-coverage modules. Based on the findings, additional unit tests were introduced for the plugin manager and policy components. In addition, ONNX unit tests were integrated into the CI workflow and successfully enabled. As a result, the overall sanity test coverage improved significantly. ### Coverage Analysis Initial coverage: * Overall coverage: **27.72%** After re-enabling eligible test cases and adding new unit tests: * Overall coverage: **40.07%** ### Coverage Improvements #### Plugin Manager Added: ```text tests/pytorch/test_plugin_manager.py ``` Coverage result: ```markdown | Improvement Area | Target File | Added Test File | Coverage After Improvement | |------------------|-------------|-----------------|----------------------------| | Core Plugin | `plugin/core/manager.py` | `tests/pytorch/test_plugin_manager.py` | **71%** | | Core Plugin | `plugin/core/policy.py` | `tests/pytorch/test_plugin_policy.py` | **99%** | | Core Backend | `plugin/core/backends/flagos/flagos.py` | `tests/pytorch/test_backend_flagos.py` | **95%** | | Core Backend Operator | `plugin/core/backends/flagos/impl/fused_adam.py` | `tests/pytorch/test_fused_adam.py` | **98%** | ``` #### ONNX Unit Tests * Added ONNX unit tests into the CI coverage workflow. * Fixed related issues and verified successful execution. * Expanded coverage of ONNX-related code paths. ### CI Test Investigation Previously excluded test groups were evaluated for re-enablement. #### Sanity * All tests can be re-enabled. #### JIT * All tests pass. #### Numerics The following tests are still failing: * `test_linear_accuracy` * `test_transformer_layer_hidden_states_format` All remaining numerics tests pass successfully. ### Changes * Added unit tests for plugin manager and policy modules. * Improved coverage of plugin core infrastructure. * Integrated and enabled ONNX unit tests in CI. * Re-evaluated previously excluded test cases and re-enabled eligible suites. * Increased overall test coverage from **27.72%** to **40.07%**. Fixes # (issue) --------- Signed-off-by: BrianPei <kaworu228@gmail.com> Co-authored-by: BrianPei <kaworu228@gmail.com> Co-authored-by: AlexMa616 <19025408700@163.com>
## Description
This PR adds FlagOS Triton implementations for fused RoPE operators that
were previously missing from the FlagOS backend.
The implementation follows the CUDA backend behavior for regular RoPE
and fused QKV RoPE paths, including forward and backward execution,
multiple QKV layouts,
interleaved and non-interleaved rotary layouts, start position offsets,
THD variable-length sequences, and context-parallel position handling.
Fixes: N/A
## Type of change
- [ ] Documentation change
- [ ] Bug fix
- [x] New feature
- [ ] Breaking change
- [ ] Infra/Build change
- [ ] Code refactoring
## Changes
- Added Triton implementations for:
- `fused_rope_forward`
- `fused_rope_backward`
- `fused_qkv_rope_forward`
- `fused_qkv_rope_backward`
- Registered the fused RoPE operators in the FlagOS backend.
- Added backend methods to expose the new FlagOS fused RoPE
implementations.
- Added tests for fused RoPE and fused QKV RoPE covering:
- `NVTE_SBHD`, `NVTE_BSHD`, and `NVTE_THD`
- interleaved and non-interleaved rotary layouts
- forward and backward paths
- start position offsets
- context-parallel position mapping
- QKV split handling, including GQA-style splits
- Updated plugin test discovery to include the fused RoPE test suite.
- Improved tests to compare FlagOS outputs against a PyTorch reference
and, when available, the CUDA vendor backend.
## Validation
- `bash ./qa/format.sh`
- `python3 transformer_engine/plugin/tests/test_fused_rope.py`
## Checklist
- [x] I have read and followed the contributing guidelines
- [x] The functionality is complete
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
…transformer_engine_hygon 2.13 (NVIDIA#85) # Description This PR fixes two compatibility issues in the Hygon backend. First, it improves the package path resolution for `transformer_engine_hygon`. In some environments, `importlib.util.find_spec()` may return a `ModuleSpec` with `origin=None`, which causes the previous implementation to fail when resolving the package directory. This PR adds a fallback to `submodule_search_locations` and reports an error if neither source is available. Second, `transformer_engine_hygon` currently does not implement the `multi_tensor_scale_tensor` API (available in upstream NVIDIA TransformerEngine v2.14). This PR replaces the unsupported call with the existing `multi_tensor_scale` implementation by extracting the scalar value from the input tensor, preserving equivalent functionality on the Hygon backend. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Improve `transformer_engine_hygon` package path resolution by handling the case where `ModuleSpec.origin` is `None`. - Add a fallback to `submodule_search_locations` when locating Hygon backend libraries. - Return a descriptive error when the package path cannot be determined. - Replace the unsupported `multi_tensor_scale_tensor` call with `multi_tensor_scale` by converting the scale tensor to a scalar value. - Add comments explaining the compatibility workaround for the Hygon backend. # Checklist: - [x] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Co-authored-by: wangyl <wangyl16@sugon.com> Co-authored-by: wangyl166 <601199939@qq.com>
- Add support for applying txda-related patches --------- Co-authored-by: malin <malin@dev.com> Co-authored-by: malin <malin@tsingmicro.com>
…g in the reference backend (NVIDIA#89) ## Summary This PR adds Ascend NPU support to the TE-FL plugin system through `torch_npu` and `transformer_engine_npu`, and fixes backward-path issues in the reference GEMM implementation. ## Changes ### Ascend NPU backend - Add automatic NPU availability detection and vendor-priority registration. - Add support for: - FlashAttention with SBHD, BSHD, and THD layouts - RMSNorm forward and backward - Generic and grouped GEMM - Multi-tensor scale and L2-norm operations - Add THD ↔ BSHD conversion operators. - Keep NPU dependencies lazily imported. ### Reference GEMM fixes - Fix output shape restoration for transposed inputs. - Do not add forward bias in backward mode. - Compute fused bias gradients. - Apply dGeLU using the saved forward activation. - Preserve correct alpha scaling and 3D input behavior. ## Testing Added coverage for: - FlashAttention forward/backward accuracy and causal masking - RMSNorm forward/backward - Generic and grouped GEMM - Multi-tensor and FP8 scale operations - Reference GEMM backward behavior Verified on Ascend 910C: ```text 46 passed ``` ## Deps It depends on TransformerEngineNPU. The package natively generated by TransformerEngineNPU is named transformer_engine. Relevant packaging logic needs to be modified so that the generated package is named transformer_engine_npu.
Move the temporary XTE TE-FL patch behavior into TE-FL native backend implementations. Register KunLunXin layernorm and GEMM operators, route attention backend selection through transformer_engine_klx_torch, and add reference GLU/DGLU fallback implementations. # Description Please include a brief summary of the changes, relevant motivation and context. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
Summary This PR adds Ascend NPU Unit CI support for TransformerEngine-FL through torch_npu and the FlagOS backend. It extends the existing CI workflow to execute real TE workloads and selected shared PyTorch tests on Ascend 910C. Changes Ascend NPU testing Add real NPU coverage for: Linear forward and backward LayerNorm, RMSNorm, and LayerNormLinear LayerNormMLP GEMM, softmax, and multi-tensor operations Unfused Dot Product Attention and MultiheadAttention TransformerLayer debug and ONNX export paths Reuse selected portable sanity and numerics tests from the existing PyTorch suites. Distributed testing Add support for: Two-process HCCL execution TE Linear gradient synchronization Context Parallel utility tests Initial non-FP8 distributed numerical validation CI and coverage Add Ascend-specific Unit test entry points. Add raw and aggregated coverage collection. Fail explicitly when torch_npu or flag_gems is unavailable. Keep unsupported CUDA-specific features explicitly excluded. Testing Verified on Ascend 910C: PyTorch Unit test execution passed. PyTorch Debug passed. PyTorch ONNX passed. PyTorch Distributed is under validation. The current PyTorch Unit job failure occurred during coverage artifact upload after the test execution had passed. Limitations CUDA Graphs, Flash/Fused Attention, FP8, MXFP8, NVFP4, block scaling, TensorRT integration, and Integration tests are not included in the current Ascend Unit scope. --------- Co-authored-by: 1395976031 <1395976031@qq.com> Co-authored-by: BrianPei <kaworu228@gmail.com> Co-authored-by: wkhylyh-debug <wkhylyh@gmail.com>
…VIDIA#92) ## Summary This PR adds a Hygon BW1000 CI baseline using the TE-FL reference backend and reorganizes the plugin tests into a backend-oriented structure under `tests/plugin`. The Hygon workflow validates the reference path only. It does not add or claim a native Hygon vendor backend. ## Changes - Add Hygon CI configuration, environment setup, and workflow entry. - Add Hygon unit, distributed smoke, ONNX smoke, and MCore integration tests. - Refactor common workflows to use platform configuration and setup scripts without chip-specific branches. - Move plugin tests from `transformer_engine/plugin/tests` to: - `tests/plugin/plugin` - `tests/plugin/backend/reference` - `tests/plugin/backend/flagos` - `tests/plugin/backend/npu` - `tests/plugin/backend/hygon` - Remove legacy plugin test files that were not collected by pytest. - Convert the FlagOS fused RoPE tests to standard pytest tests. - Add documentation for adding and running tests locally. - Use the unified 8-GPU runner labels. ## Hygon Baseline - Hardware: Hygon BW1000 - Backend policy: `TE_FL_PREFER=reference` - GEMM implementation: `reference.torch` - Runner label: `hg-8g-cicd-te` - Coverage enabled but not required - Debug tests are explicitly skipped when `nvdlfw_inspect` is unavailable ## Testing Validated on Hygon BW1000: - PyTorch unit tests - Plugin manager and policy tests - Reference backend tests - Distributed smoke tests - ONNX smoke tests - Coverage aggregation - Megatron-LM-FL MCore integration test All configured Hygon CI jobs passed. --------- Co-authored-by: wkhylyh-debug <wkhylyh@gmail.com>
## Summary Add a dedicated MUSA CI workflow for TransformerEngine-FL. ## Changes - Added MUSA hardware configuration and workflow entry points. - Added MUSA environment setup and runtime verification. - Verified the availability of the `transformer_engine_musa_torch` API and `vendor.musa`. - Verified that the representative `generic_gemm` dispatch selects `vendor.musa`. - Executed supported native TE tests using a MUSA-specific launcher. - Added a dedicated launcher for MUSA Megatron-LM integration tests. - Increased the timeout for shared unit tests from 60 minutes to 180 minutes. ## Test Organization - Native TE test adapter: `tests/plugin/backend/musa/run_native_tests.sh` - MUSA MCore integration test entry point: `tests/integration/musa/run_mcore.sh` Unsupported MUSA test cases are filtered out within the platform-specific launcher to ensure MUSA compatibility. ## Verification - YAML configuration parsing passed. - Bash syntax checks passed. - Python setup script checks passed. - `git diff --check` passed. - The final branch has been synchronized with `origin/musa-dev`. The MUSA backend implementation already exists in the upstream source code; this change provides the corresponding dedicated CI setup and test entry points. --------- Co-authored-by: canghaiX <1395976031@qq.com> Co-authored-by: BrianPei <kaworu228@gmail.com> Co-authored-by: wkhylyh-debug <wkhylyh@gmail.com> Co-authored-by: canghaiX <59075364+canghaiX@users.noreply.github.com>
# Conflicts: # .github/actions/build-pytorch-wheel/Dockerfile # .github/workflows/build.yml # .github/workflows/lint.yml # .gitignore # README.rst # benchmarks/linear/benchmark_grouped_linear.py # build_tools/VERSION.txt # build_tools/wheel_utils/Dockerfile.aarch # build_tools/wheel_utils/Dockerfile.x86 # build_tools/wheel_utils/build_wheels.sh # docs/api/pytorch.rst # docs/conf.py # docs/debug/3_api_features.rst # docs/envvars.rst # docs/examples/op_fuser/op_fuser.rst # docs/features/low_precision_training/index.rst # docs/index.rst # examples/jax/collective_gemm/common.py # examples/jax/collective_gemm/run_test_cgemm.sh # examples/jax/collective_gemm/test_gemm.py # examples/pytorch/quantized_model_init/fully_shard.py # qa/L0_jax_unittest/test.sh # qa/L0_pytorch_debug_unittest/test.sh # qa/L0_pytorch_unittest/test.sh # qa/L1_jax_distributed_unittest/test.sh # qa/L1_pytorch_distributed_unittest/test.sh # qa/L1_pytorch_thunder_integration/test.sh # qa/L2_jax_distributed_unittest/test.sh # qa/L3_pytorch_FA_versions_test/test.sh # setup.py # tests/cpp/operator/test_cast_mxfp8_grouped.cu # tests/cpp/operator/test_cast_nvfp4_transpose.cu # tests/cpp/operator/test_grouped_gemm.cu # tests/cpp/operator/test_swizzle.cu # tests/cpp/test_common.cu # tests/cpp/test_common.h # tests/jax/test_custom_call_compute.py # tests/jax/test_fused_attn.py # tests/jax/test_fused_router.py # tests/pytorch/attention/run_attention_with_cp.py # tests/pytorch/attention/test_attention.py # tests/pytorch/attention/test_attention_with_cp.py # tests/pytorch/debug/test_api_features.py # tests/pytorch/debug/test_log.py # tests/pytorch/debug/test_perf.py # tests/pytorch/debug/test_sanity.py # tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py # tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py # tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py # tests/pytorch/distributed/run_numerics_exact.py # tests/pytorch/distributed/test_cast_master_weights_to_fp8.py # tests/pytorch/distributed/test_comm_gemm_overlap.py # tests/pytorch/distributed/test_torch_fsdp2.py # tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py # tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py # tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py # tests/pytorch/nvfp4/test_nvfp4_group_quantize.py # tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py # tests/pytorch/nvfp4/test_nvfp4_module_exact.py # tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py # tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py # tests/pytorch/test_cpu_offloading.py # tests/pytorch/test_custom_recipe.py # tests/pytorch/test_float8_current_scaling_exact.py # tests/pytorch/test_float8blockwisetensor.py # tests/pytorch/test_fused_router.py # tests/pytorch/test_fusible_ops.py # tests/pytorch/test_grouped_tensor.py # tests/pytorch/test_numerics.py # tests/pytorch/test_onnx_export.py # tests/pytorch/test_permutation.py # tests/pytorch/test_quantized_tensor.py # tests/pytorch/test_recipe.py # tests/pytorch/test_sanity.py # tests/pytorch/utils.py # transformer_engine/__init__.py # transformer_engine/common/CMakeLists.txt # transformer_engine/common/__init__.py # transformer_engine/common/cast/cast.cu # transformer_engine/common/cast/core/common.cuh # transformer_engine/common/cast/dispatch/dequantize.cuh # transformer_engine/common/cast/dispatch/quantize.cuh # transformer_engine/common/cast/fp8/gated_fp8.cuh # transformer_engine/common/cast/fp8/quantize_fp8.cuh # transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh # transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh # transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh # transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh # transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh # transformer_engine/common/cast/mxfp8/swizzle.cuh # transformer_engine/common/cast/nvfp4/core_nvfp4.cuh # transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh # transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh # transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh # transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh # transformer_engine/common/comm_gemm/comm_gemm.cpp # transformer_engine/common/common.cu # transformer_engine/common/common.h # transformer_engine/common/fused_attn/fused_attn.cpp # transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu # transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h # transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu # transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h # transformer_engine/common/fused_attn/fused_attn_fp8.cu # transformer_engine/common/fused_attn/fused_attn_fp8.h # transformer_engine/common/fused_attn/utils.h # transformer_engine/common/fused_router/fused_moe_aux_loss.cu # transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu # transformer_engine/common/fused_router/fused_topk_with_score_function.cu # transformer_engine/common/fused_router/utils.h # transformer_engine/common/gemm/cublaslt_grouped_gemm.cu # transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu # transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu # transformer_engine/common/hadamard_transform/group_hadamard_transform.cu # transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu # transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu # transformer_engine/common/hadamard_transform/hadamard_transform.cu # transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu # transformer_engine/common/include/transformer_engine/fused_attn.h # transformer_engine/common/include/transformer_engine/gemm.h # transformer_engine/common/include/transformer_engine/swizzle.h # transformer_engine/common/include/transformer_engine/transformer_engine.h # transformer_engine/common/include/transformer_engine/utils.h # transformer_engine/common/multi_tensor/compute_scale.cu # transformer_engine/common/multi_tensor/scale.cu # transformer_engine/common/normalization/common.cpp # transformer_engine/common/recipe/__init__.py # transformer_engine/common/recipe/nvfp4.cu # transformer_engine/common/swizzle/swizzle.cu # transformer_engine/common/transformer_engine.cpp # transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu # transformer_engine/common/triton/permutation.py # transformer_engine/common/util/utils.cu # transformer_engine/debug/features/api.py # transformer_engine/debug/features/fake_quant.py # transformer_engine/debug/features/log_fp8_tensor_stats.py # transformer_engine/debug/features/per_tensor_scaling.py # transformer_engine/jax/attention.py # transformer_engine/jax/cpp_extensions/__init__.py # transformer_engine/jax/cpp_extensions/gemm.py # transformer_engine/jax/cpp_extensions/quantization.py # transformer_engine/jax/cpp_extensions/router.py # transformer_engine/jax/csrc/extensions.h # transformer_engine/jax/csrc/extensions/attention.cpp # transformer_engine/jax/csrc/extensions/gemm.cpp # transformer_engine/jax/csrc/extensions/inspect.cpp # transformer_engine/jax/csrc/extensions/misc.h # transformer_engine/jax/csrc/extensions/pybind.cpp # transformer_engine/jax/csrc/extensions/quantization.cpp # transformer_engine/jax/csrc/extensions/router.cpp # transformer_engine/jax/debug/experimental/inspect.py # transformer_engine/jax/dense.py # transformer_engine/jax/flax/__init__.py # transformer_engine/jax/flax/module.py # transformer_engine/jax/flax/transformer.py # transformer_engine/jax/permutation.py # transformer_engine/jax/quantize/tensor.py # transformer_engine/jax/router.py # transformer_engine/jax/triton_extensions/permutation.py # transformer_engine/jax/triton_extensions/utils.py # transformer_engine/jax/version_utils.py # transformer_engine/pytorch/__init__.py # transformer_engine/pytorch/attention/dot_product_attention/backends.py # transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py # transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py # transformer_engine/pytorch/attention/dot_product_attention/utils.py # transformer_engine/pytorch/attention/multi_head_attention.py # transformer_engine/pytorch/constants.py # transformer_engine/pytorch/cpp_extensions/fused_attn.py # transformer_engine/pytorch/cpp_extensions/gemm.py # transformer_engine/pytorch/cpu_offload.py # transformer_engine/pytorch/cpu_offload_v1.py # transformer_engine/pytorch/csrc/common.h # transformer_engine/pytorch/csrc/extensions.h # transformer_engine/pytorch/csrc/extensions/attention.cpp # transformer_engine/pytorch/csrc/extensions/cast.cpp # transformer_engine/pytorch/csrc/extensions/gemm.cpp # transformer_engine/pytorch/csrc/extensions/misc.cpp # transformer_engine/pytorch/csrc/extensions/pybind.cpp # transformer_engine/pytorch/csrc/extensions/router.cpp # transformer_engine/pytorch/csrc/extensions/swizzle.cpp # transformer_engine/pytorch/csrc/quantizer.cpp # transformer_engine/pytorch/csrc/type_converters.cpp # transformer_engine/pytorch/csrc/util.h # transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py # transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py # transformer_engine/pytorch/distributed.py # transformer_engine/pytorch/experimental/quantization_nvfp4.py # transformer_engine/pytorch/jit.py # transformer_engine/pytorch/module/base.py # transformer_engine/pytorch/module/grouped_linear.py # transformer_engine/pytorch/module/layernorm_linear.py # transformer_engine/pytorch/module/layernorm_mlp.py # transformer_engine/pytorch/module/linear.py # transformer_engine/pytorch/ops/__init__.py # transformer_engine/pytorch/ops/_common.py # transformer_engine/pytorch/ops/basic/__init__.py # transformer_engine/pytorch/ops/basic/activation.py # transformer_engine/pytorch/ops/basic/basic_linear.py # transformer_engine/pytorch/ops/basic/bias.py # transformer_engine/pytorch/ops/basic/grouped_linear.py # transformer_engine/pytorch/ops/basic/swiglu.py # transformer_engine/pytorch/ops/fused/__init__.py # transformer_engine/pytorch/ops/fused/backward_activation_bias.py # transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py # transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py # transformer_engine/pytorch/ops/fuser.py # transformer_engine/pytorch/optimizers/fused_adam.py # transformer_engine/pytorch/permutation.py # transformer_engine/pytorch/quantization.py # transformer_engine/pytorch/quantized_tensor.py # transformer_engine/pytorch/router.py # transformer_engine/pytorch/tensor/_quantization_helpers.py # transformer_engine/pytorch/tensor/float8_blockwise_tensor.py # transformer_engine/pytorch/tensor/float8_tensor.py # transformer_engine/pytorch/tensor/grouped_tensor.py # transformer_engine/pytorch/tensor/mxfp8_tensor.py # transformer_engine/pytorch/tensor/nvfp4_tensor.py # transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py # transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py # transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py # transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py # transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py # transformer_engine/pytorch/tensor/utils.py # transformer_engine/pytorch/transformer.py # transformer_engine/pytorch/triton/__init__.py # transformer_engine/pytorch/triton/permutation.py # transformer_engine/pytorch/utils.py
|
wendell seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
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
19 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: