Split agent guidance into task-specific documents

AGENTS.md

@@ -127,219 +127,29 @@ frontend-only change does not require validation of the other frontend, though
 shared behavior, documentation, serialization, and feature parity may still be
 affected.
 
-## Build and Environment
-
-Use a compatible NGC PyTorch or JAX development container when possible. For a
-host build, follow `docs/installation.rst`. The selected framework must be
-installed before Transformer Engine because the root build imports it while
-configuring the framework extension. Development builds therefore use
-`--no-build-isolation`.
-
-Initialize the pinned source dependencies before the first build:
-
-```bash
-git submodule update --init --recursive
-```
-
-The root editable build is the normal development build. Run it from the
-repository root. The frontend is detected automatically but can also be selected explicitly
-via NVTE_FRAMEWORK environment variable during build:
-
-```bash
-NVTE_FRAMEWORK=pytorch python -m pip install -e . -v --no-build-isolation
-NVTE_FRAMEWORK=jax python -m pip install -e . -v --no-build-isolation
-NVTE_FRAMEWORK=pytorch,jax python -m pip install -e . -v --no-build-isolation
-NVTE_FRAMEWORK=none python -m pip install -e . -v --no-build-isolation
-```
-
-Common build controls are:
-
-- `NVTE_CUDA_ARCHS`: semicolon-separated target compute capabilities. For
-  development, set this to the exact architecture being tested to avoid
-  compiling unused variants. On Blackwell, use the architecture-specific
-  target such as `100a` rather than the `100` family selector, which expands
-  to multiple variants. Use a family selector or explicit list when building
-  for multiple GPU architectures or producing portable artifacts.
-- `MAX_JOBS` or `NVTE_BUILD_MAX_JOBS`: concurrent compilation jobs.
-- `NVTE_BUILD_THREADS_PER_JOB`: threads within a build job.
-- `NVTE_USE_CCACHE=1` and `NVTE_CCACHE_BIN`: compiler caching.
-- `NVTE_CMAKE_BUILD_DIR`: alternate common-library build directory.
-- `NVTE_BUILD_DEBUG=1`: debug build.
-- `CUDA_HOME`, `CUDA_PATH`, `CUDNN_PATH`, and `CXX`: toolchain selection.
-
-Before limiting compilation to a small fixed number of jobs or threads, check
-the number of CPU cores available on the system and choose concurrency that
-fits the current machine.
-
-Optional native components add discovery and linking requirements:
-
-- `NVTE_WITH_NCCL_EP`: NCCL expert parallelism; enabled by default for
-  applicable Hopper-or-newer targets and requires compatible NCCL headers and
-  libraries.
-- `NVTE_UB_WITH_MPI=1`: MPI userbuffers bootstrap; requires `MPI_HOME`.
-- `NVTE_ENABLE_NVSHMEM=1`: NVSHMEM; requires `NVSHMEM_HOME`.
-- `NVTE_WITH_CUBLASMP=1`: cuBLASMp; may require `CUBLASMP_HOME`.
-- `NVTE_WITH_CUSOLVERMP=1`: cuSolverMp; may require `CUSOLVERMP_HOME`.
-
-See `docs/envvars.rst`, `setup.py`, `build_tools/`, and
-`transformer_engine/common/CMakeLists.txt` for less common settings.
-
-### Rebuild scope
-
-- Python-only change: no rebuild after an editable install; restart processes
-  that already imported the module.
-- C++ change, either in common code or a framework binding: rerun the same build
-  command and configuration used to build Transformer Engine before the change.
-  This preserves the selected frontends, architecture targets, and optional
-  features.
-- Build logic, compiler flags, optional features, or framework selection:
-  rerun the full editable build; use a clean build tree if its generated state
-  is incompatible.
-- Submodule source or revision: rebuild native targets and validate the
-  dependent component.
-
-CMake normally reuses `build/cmake`. Before removing a build tree, check
-`NVTE_CMAKE_BUILD_DIR` and the verbose build output. A clean tree is commonly
-needed after changing toolchains, CMake generators, architecture targets, or
-optional native dependencies.
-
-### Verify the imported build
-
-NGC images may already contain Transformer Engine. After building, confirm that
-the framework sees a GPU and that TE resolves to the current worktree:
-
-Sandboxing may block access to GPU devices. If a GPU visibility check fails or
-reports no devices, do not conclude that the system has no GPUs until the check
-has also been attempted outside the sandbox with the required permission.
-
-```bash
-python -c 'import torch, transformer_engine, transformer_engine.pytorch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available()); print(transformer_engine.__file__)'
-```
-
-```bash
-python -c 'import jax, transformer_engine, transformer_engine.jax; print(jax.__version__, jax.devices()); print(transformer_engine.__file__)'
-```
-
-## Testing
-
-PyTorch and JAX tests require `pytest`. The `qa/**/test.sh` launchers install
-the pytest version and plugins used by their suites and contain the
-authoritative environment variables, configuration files, and command-line
-options for specialized tests. Before running a test directly, find its
-invocation or the closest analogous test in the relevant launcher and preserve
-those settings. Use a focused `python -m pytest` command during iteration, then
-use the launcher for the corresponding full suite.
-
-### Focused framework tests
-
-```bash
-# PyTorch
-python -m pytest -v tests/pytorch/path/to/test_file.py -k 'specific_case'
-python -m pytest -v tests/pytorch/test_sanity.py
-
-# JAX: use the repository pytest configuration
-python -m pytest -c tests/jax/pytest.ini -v \
-    tests/jax/path/to/test_file.py -k 'specific_case'
-python -m pytest -c tests/jax/pytest.ini -v tests/jax/test_layer.py
-```
-
-### Common C++ and CUDA tests
-
-Separately built after the main build.
-
-```bash
-cmake -GNinja -S tests/cpp -B tests/cpp/build \
-    -DCMAKE_CUDA_ARCHITECTURES="$NVTE_CUDA_ARCHS"
-cmake --build tests/cpp/build
-ctest --test-dir tests/cpp/build --output-on-failure
-```
-
-For a focused GoogleTest run:
-
-```bash
-tests/cpp/build/operator/test_operator --gtest_list_tests
-tests/cpp/build/operator/test_operator --gtest_filter='*RelevantPattern*'
-```
-
-### QA launchers
-
-`CONTRIBUTING.rst` requires the `L0_*` checks. Their scripts default to
-`TE_PATH=/opt/transformerengine` and `XML_LOG_DIR=/logs`; override both for a
-local checkout:
-
-```bash
-TE_PATH="$PWD" XML_LOG_DIR=/tmp/te-test-logs \
-    bash qa/L0_pytorch_unittest/test.sh
-TE_PATH="$PWD" XML_LOG_DIR=/tmp/te-test-logs \
-    bash qa/L0_jax_unittest/test.sh
-TE_PATH="$PWD" XML_LOG_DIR=/tmp/te-test-logs \
-    bash qa/L0_cppunittest/test.sh
-```
-
-Inspect a QA script before running it: some install dependencies or require a
-particular container, GPU architecture, GPU count, or external repository.
-Use the higher-level launcher for distributed, Megatron Core, FSDP, ONNX, and
-attention-backend coverage rather than reconstructing its environment.
-
-### Validation by changed area
-
-- Shared Python: focused behavior through each affected frontend.
-- PyTorch or JAX API: focused frontend test plus related integration coverage.
-- Private framework binding: loaded-extension check and a focused numerical or
-  integration test.
-- Common C/CUDA: focused C++ operator test plus each affected frontend.
-- Build or packaging: build in an appropriate clean environment, verify
-  artifact paths, then use `qa/L0_pytorch_wheel/` or `qa/L0_jax_wheel/`.
-- Distributed component: smallest supported multi-process case, then its QA
-  launcher; record the GPU count and topology.
-- Public API: tests, documentation, and compatibility across supported callers.
-
-## Numerical and Backend Validation
-
-For numerical or low-precision changes:
-
-- Compare outputs and applicable gradients with an independent native-framework
-  or higher-precision reference, not only shapes or successful execution.
-- Cover the affected dtypes, recipes, layouts, boundary shapes, forward and
-  backward paths, and GPU architectures.
-- Base tolerances on the operation, dtype, accumulation, and reference. Do not
-  loosen them solely to make a failing case pass.
-- Record flags such as `NVTE_ALLOW_NONDETERMINISTIC_ALGO` that change expected
-  numerical behavior.
-- Exercise the intended backend and fallback behavior when the operation can
-  use cuDNN, FlashAttention, unfused attention, Triton, NVRTC, or an
-  architecture-specific kernel.
-- Give skips and expected failures a capability-based or tracked-defect reason;
-  do not remove a failing configuration before identifying its cause.
-
-Choose tolerances for the quantity being compared, not merely for the lowest
-precision used internally. For example, an MXFP8 GEMM that returns BF16 should
-use BF16-appropriate output tolerances rather than FP8-appropriate tolerances
-when its operands are chosen to be exactly representable in MXFP8. One way to
-construct such operands is to generate random values, quantize and dequantize
-them, and use the resulting values in both the reference and operation under
-test. Apply this construction to every quantized operand, including parameters
-held by a module such as weights, not only to the module's explicit inputs.
-
-When reporting a GPU or distributed failure, include the exact command, input
-shape/dtype/layout, relevant `NVTE_*` variables, GPU model and count, compute
-capability, topology, and framework/CUDA/cuDNN/NCCL/driver versions. Also verify
-which backend was selected and that the imported extension came from the
-current worktree. Reduce failures across ranks, shapes, recipes, and backends
-while preserving the observed behavior.
-
-## Performance Changes
-
-Use the relevant program under `benchmarks/` and report its exact command. A
-performance result must identify the selected TE backend or kernel, GPU model
-and count, software versions, shapes, layouts, dtypes, recipes, warm-up,
-synchronization, input distribution, and measurement statistic. Separate
-compile or autotune time, Python launch overhead, communication, and
-steady-state kernel time where they materially affect the result.
-
-Compare on equivalent hardware and software configurations and confirm
-numerical correctness before accepting a faster result. Do not place one-off
-performance measurements in correctness tests.
+## Task-specific guidance
+
+Read the guides relevant to the work being performed. Paths and commands in
+these guides are relative to the repository root. A task may require
+more than one guide but be wary of using more context than needed.
+
+- [Building](docs/development/building.md): read when installing or rebuilding
+  Transformer Engine, configuring its development environment, changing
+  toolchains or native build options, or verifying the imported installation.
+  Covers framework/container selection, source dependencies, editable
+  installation, architecture targets, build controls, and rebuild scope.
+
+- [Testing](docs/development/testing.md): read when preparing the test
+  environment, choosing or running tests, reproducing CI test failures,
+  or validating numerical or backend behavior.
+  Covers test dependencies such as pytest and plugins, suite-specific
+  environment variables and options, QA launchers, C++ test executables,
+  and frontend, numerical, backend, and distributed validation.
+
+- [Performance](docs/development/performance.md): read when benchmarking,
+  investigating performance changes, or reporting performance comparisons.
+  Covers benchmark selection, measurement conditions, correctness checks,
+  and the information required to reproduce results.
 
 ## Formatting and Linting
 

docs/development/building.md

@@ -0,0 +1,92 @@
+## Build and Environment
+
+Use a compatible NGC PyTorch or JAX development container when possible. For a
+host build, follow `docs/installation.rst`. The selected framework must be
+installed before Transformer Engine because the root build imports it while
+configuring the framework extension. Development builds therefore use
+`--no-build-isolation`.
+
+Initialize the pinned source dependencies before the first build:
+
+```bash
+git submodule update --init --recursive
+```
+
+The root editable build is the normal development build. Run it from the
+repository root. The frontend is detected automatically but can also be selected explicitly
+via NVTE_FRAMEWORK environment variable during build:
+
+```bash
+NVTE_FRAMEWORK=pytorch python -m pip install -e . -v --no-build-isolation
+NVTE_FRAMEWORK=jax python -m pip install -e . -v --no-build-isolation
+NVTE_FRAMEWORK=pytorch,jax python -m pip install -e . -v --no-build-isolation
+NVTE_FRAMEWORK=none python -m pip install -e . -v --no-build-isolation
+```
+
+Common build controls are:
+
+- `NVTE_CUDA_ARCHS`: semicolon-separated target compute capabilities. For
+  development, set this to the exact architecture being tested to avoid
+  compiling unused variants. On Blackwell, use the architecture-specific
+  target such as `100a` rather than the `100` family selector, which expands
+  to multiple variants. Use a family selector or explicit list when building
+  for multiple GPU architectures or producing portable artifacts.
+- `MAX_JOBS` or `NVTE_BUILD_MAX_JOBS`: concurrent compilation jobs.
+- `NVTE_BUILD_THREADS_PER_JOB`: threads within a build job.
+- `NVTE_USE_CCACHE=1` and `NVTE_CCACHE_BIN`: compiler caching.
+- `NVTE_CMAKE_BUILD_DIR`: alternate common-library build directory.
+- `NVTE_BUILD_DEBUG=1`: debug build.
+- `CUDA_HOME`, `CUDA_PATH`, `CUDNN_PATH`, and `CXX`: toolchain selection.
+
+Before limiting compilation to a small fixed number of jobs or threads, check
+the number of CPU cores available on the system and choose concurrency that
+fits the current machine.
+
+Optional native components add discovery and linking requirements:
+
+- `NVTE_WITH_NCCL_EP`: NCCL expert parallelism; enabled by default for
+  applicable Hopper-or-newer targets and requires compatible NCCL headers and
+  libraries.
+- `NVTE_UB_WITH_MPI=1`: MPI userbuffers bootstrap; requires `MPI_HOME`.
+- `NVTE_ENABLE_NVSHMEM=1`: NVSHMEM; requires `NVSHMEM_HOME`.
+- `NVTE_WITH_CUBLASMP=1`: cuBLASMp; may require `CUBLASMP_HOME`.
+- `NVTE_WITH_CUSOLVERMP=1`: cuSolverMp; may require `CUSOLVERMP_HOME`.
+
+See `docs/envvars.rst`, `setup.py`, `build_tools/`, and
+`transformer_engine/common/CMakeLists.txt` for less common settings.
+
+### Rebuild scope
+
+- Python-only change: no rebuild after an editable install; restart processes
+  that already imported the module.
+- C++ change, either in common code or a framework binding: rerun the same build
+  command and configuration used to build Transformer Engine before the change.
+  This preserves the selected frontends, architecture targets, and optional
+  features.
+- Build logic, compiler flags, optional features, or framework selection:
+  rerun the full editable build; use a clean build tree if its generated state
+  is incompatible.
+- Submodule source or revision: rebuild native targets and validate the
+  dependent component.
+
+CMake normally reuses `build/cmake`. Before removing a build tree, check
+`NVTE_CMAKE_BUILD_DIR` and the verbose build output. A clean tree is commonly
+needed after changing toolchains, CMake generators, architecture targets, or
+optional native dependencies.
+
+### Verify the imported build
+
+NGC images may already contain Transformer Engine. After building, confirm that
+the framework sees a GPU and that TE resolves to the current worktree:
+
+Sandboxing may block access to GPU devices. If a GPU visibility check fails or
+reports no devices, do not conclude that the system has no GPUs until the check
+has also been attempted outside the sandbox with the required permission.
+
+```bash
+python -c 'import torch, transformer_engine, transformer_engine.pytorch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available()); print(transformer_engine.__file__)'
+```
+
+```bash
+python -c 'import jax, transformer_engine, transformer_engine.jax; print(jax.__version__, jax.devices()); print(transformer_engine.__file__)'
+```

docs/development/performance.md

@@ -0,0 +1,12 @@
+## Performance Changes
+
+Use the relevant program under `benchmarks/` and report its exact command. A
+performance result must identify the selected TE backend or kernel, GPU model
+and count, software versions, shapes, layouts, dtypes, recipes, warm-up,
+synchronization, input distribution, and measurement statistic. Separate
+compile or autotune time, Python launch overhead, communication, and
+steady-state kernel time where they materially affect the result.
+
+Compare on equivalent hardware and software configurations and confirm
+numerical correctness before accepting a faster result. Do not place one-off
+performance measurements in correctness tests.