From c8d1bbd595264efe834cb12e92a2e1ac1db8c530 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 30 Jun 2026 15:15:10 -0500 Subject: [PATCH 01/33] Add optional JAX support to FE OSS APIs --- docs/fe-oss-apis/cutedsl-jax-design.md | 754 ++++++++++++++++++ docs/fe-oss-apis/overview.md | 28 +- docs/fe-oss-apis/rmsnorm_rht_amax.md | 45 +- pyproject.toml | 5 + python/cudnn/README.md | 52 +- python/cudnn/_rmsnorm_rht_amax_config.py | 92 +++ python/cudnn/frontend/__init__.py | 47 ++ python/cudnn/frontend/_legacy_gaps.py | 159 ++++ python/cudnn/frontend/_registry.py | 319 ++++++++ python/cudnn/frontend/rmsnorm_rht_amax.py | 69 ++ python/cudnn/jax/__init__.py | 8 + python/cudnn/jax/cutedsl.py | 484 +++++++++++ python/cudnn/jax/rmsnorm_rht_amax.py | 151 ++++ python/cudnn/rmsnorm_rht_amax/__init__.py | 24 +- .../fe_api/test_frontend_target_parity.py | 551 +++++++++++++ .../python/fe_api/test_jax_cutedsl_adapter.py | 365 +++++++++ .../fe_api/test_jax_rmsnorm_rht_amax.py | 110 +++ .../fe_api/test_rmsnorm_rht_amax_config.py | 84 ++ 18 files changed, 3332 insertions(+), 15 deletions(-) create mode 100644 docs/fe-oss-apis/cutedsl-jax-design.md create mode 100644 python/cudnn/_rmsnorm_rht_amax_config.py create mode 100644 python/cudnn/frontend/__init__.py create mode 100644 python/cudnn/frontend/_legacy_gaps.py create mode 100644 python/cudnn/frontend/_registry.py create mode 100644 python/cudnn/frontend/rmsnorm_rht_amax.py create mode 100644 python/cudnn/jax/__init__.py create mode 100644 python/cudnn/jax/cutedsl.py create mode 100644 python/cudnn/jax/rmsnorm_rht_amax.py create mode 100644 test/python/fe_api/test_frontend_target_parity.py create mode 100644 test/python/fe_api/test_jax_cutedsl_adapter.py create mode 100644 test/python/fe_api/test_jax_rmsnorm_rht_amax.py create mode 100644 test/python/fe_api/test_rmsnorm_rht_amax_config.py diff --git a/docs/fe-oss-apis/cutedsl-jax-design.md b/docs/fe-oss-apis/cutedsl-jax-design.md new file mode 100644 index 000000000..a38f4edd5 --- /dev/null +++ b/docs/fe-oss-apis/cutedsl-jax-design.md @@ -0,0 +1,754 @@ +# CuTe DSL + JAX support for FE-OSS APIs + +Status: design proposal with a proof of concept. The public API and internal +interfaces described here are experimental. + +## Executive decision + +Do not make the existing `APIBase.compile()` / `execute()` lifecycle accept JAX +arrays, and do not introduce a second replacement Torch API. The existing +top-level Torch classes and wrappers remain canonical and backward compatible. +Add JAX as an optional functional namespace under `cudnn.jax`, lowering CuTe DSL +launchers through `cutlass.jax.cutlass_call`. + +An internal framework-neutral catalog connects those intentionally different +public APIs. It records semantic parameter/output mappings, exact API and kernel +ownership, and whether JAX is implemented or is an explicit tracked gap. + +The distinction matters because a JAX value encountered during tracing is an +abstract value, not a device buffer: + +| Concern | Current Torch FE-OSS binding | JAX binding | +| --- | --- | --- | +| Inputs | Concrete `torch.Tensor` objects | JAX arrays or tracers with abstract shape/dtype | +| Outputs | Allocated in Python with `torch.empty*` | Declared with `ShapeDtypeStruct`; XLA allocates them | +| Kernel compilation | Explicit `cute.compile()` before execution | CuTe compilation during XLA lowering inside `cutlass_call` | +| Execution | Python invokes a compiled callable | XLA invokes an embedded custom call without a Python callback | +| CUDA stream | Caller/default Torch or CUTLASS stream | Runtime stream supplied by XLA/PJRT | +| Workspace | Often a mutable tensor cached on the API object | Per-invocation XLA-owned temporary result | +| Mutation | Preallocated outputs are mutated | Public API remains functional; legal aliases are declared | +| Cache | Python object and compiled callable | JAX executable cache plus CUTLASS function/spec cache | + +The proof of concept implements this split and ports one real FE-OSS operation, +RMSNorm + RHT + amax, without changing its existing Torch API. It also adds a +dependency-free ownership audit so a new public API or physical CuTe kernel +cannot silently omit its JAX support status. + +## Scope + +### Goals + +- Run FE-OSS CuTe DSL kernels inside `jax.jit` without host callbacks. +- Preserve XLA's stream ordering, allocation, liveness, and functional + semantics. +- Share operator configuration, shape inference, and kernel source between + Torch and JAX. +- Model workspace, output initialization, and legal input/output aliases. +- Establish explicit policies for autodiff, batching, sharding, dynamic shapes, + CUDA graphs, and caching. +- Avoid making Torch a dependency of the JAX-only path. + +### Non-goals for the first production increment + +- Making the existing class API traceable. +- Supporting every FE-OSS operator immediately. +- Arbitrary Torch-style strided JAX views. +- Data-dependent output or workspace sizes. +- Automatic gradients, `vmap`, or automatic SPMD partitioning for every + operator. +- A JAX binding for the C++ cuDNN graph API. That is a related but separate + adapter described later. + +## Sources reviewed + +The review started from the published +[Frontend OSS API guide](https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/fe-oss-apis.html) +and covered the local implementation under `python/cudnn`, its tests, and +representative operators with and without workspace. The main coupling points +are: + +- [`api_base.py`](../../python/cudnn/api_base.py): `TensorDesc` stores + `torch.dtype` and `torch.device`; dtype checks and tensor conversion are + Torch-specific. +- [`datatypes.py`](../../python/cudnn/datatypes.py): framework mappings only + cover Torch today. +- [`gemm_amax/api.py`](../../python/cudnn/gemm_amax/api.py) and + [`sdpa/fwd/api.py`](../../python/cudnn/sdpa/fwd/api.py): compile fake CuTe + tensors with a fake stream, then execute against preallocated Torch tensors. +- [`sdpa/bwd/api.py`](../../python/cudnn/sdpa/bwd/api.py) and grouped GEMM APIs: + retain mutable workspaces on cached API objects. +- [`rmsnorm_rht_amax/api.py`](../../python/cudnn/rmsnorm_rht_amax/api.py): a + small, single-launch operator suitable for the first vertical slice. + +The CUTLASS package pinned by this repository is `nvidia-cutlass-dsl==4.5.0`. +The review used tag +[`v4.5.0`](https://github.com/NVIDIA/cutlass/tree/v4.5.0) at commit +`e406c186f510a15091cce01f782020ceb7ba8eb5`, plus current `main` at commit +[`e8ecfad`](https://github.com/NVIDIA/cutlass/commit/e8ecfad75b44d1ad56264f5001d877e9e47fe080) +to identify subsequent API changes. Relevant sources are: + +- [`cutlass.jax.cutlass_call`](https://github.com/NVIDIA/cutlass/blob/v4.5.0/python/CuTeDSL/cutlass/jax/primitive.py) +- [CuTe compilation and cache](https://github.com/NVIDIA/cutlass/blob/v4.5.0/python/CuTeDSL/cutlass/jax/compile.py) +- [JAX tensor/layout adapter](https://github.com/NVIDIA/cutlass/blob/v4.5.0/python/CuTeDSL/cutlass/jax/types.py) +- [basic, layout, and alias examples](https://github.com/NVIDIA/cutlass/blob/v4.5.0/examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_basic.py) +- [symbolic export example](https://github.com/NVIDIA/cutlass/blob/v4.5.0/examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_export.py) +- [`shard_map` and custom partitioning example](https://github.com/NVIDIA/cutlass/blob/v4.5.0/examples/python/CuTeDSL/dsl_tutorials/jax/cutlass_call_sharding.py) + +The `python/CuTeDSL/cutlass/jax` implementation files carry NVIDIA's +`LicenseRef-NvidiaProprietary` header even though they are visible in the public +repository. They are architecture references only; no source is copied into +this OSS proof of concept. CUTLASS's example files use BSD-3-Clause. + +The JAX-side behavior was checked against the official +[CuTe DSL + JAX guide](https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html), +[FFI guide](https://docs.jax.dev/en/latest/ffi.html), and +[`jax.ffi.ffi_call` API](https://docs.jax.dev/en/latest/_autosummary/jax.ffi.ffi_call.html). + +## What `cutlass_call` provides + +`cutlass_call` is the correct CuTe DSL integration seam. It already supplies the +custom JAX primitive needed to compile Python CuTe code during lowering: + +1. Normal JAX tracing flattens input pytrees and records output abstract values. +2. The CUDA lowering builds a CUTLASS function specification from abstract + shapes, dtypes, layouts, aliases, compile flags, and static keyword arguments. +3. It calls `cute.compile` with fake JAX descriptors and a placeholder stream. +4. It emits a position-independent object containing the host launch wrapper and + device code, then embeds those bytes in a typed XLA FFI custom call. +5. At execution, the CUTLASS runtime invokes that object with XLA-owned buffers + and XLA's CUDA stream. Python is not on the execution path. + +The placeholder stream used during compilation is not the runtime stream. FE +launchers must accept the stream as their first argument and pass it to every +helper and main-kernel launch. + +The bridge supports explicit output metadata, compact physical layouts, logical +mode permutations, input/output aliases, static or runtime tensor dimensions, +CUDA-graph opt-out, and `jax.export`. It deliberately rejects autodiff and +`vmap`; sharding needs `shard_map` or an operator-specific partitioner. + +## Proposed architecture + +### Parity by construction + +The unit of API parity is a **semantic operation variant**, not an individual +`@cute.kernel`. A forward operation may own a main kernel, initialization +helpers, and architecture-specific implementations while still presenting one +Torch/JAX contract. Conversely, a public helper that is independently callable +may deserve its own operation entry even when it shares a kernel module. + +Each operation is registered once with: + +```text +FrontendOperationSpec + name stable semantic/variant identifier + contract_signature internal common parameters/defaults + targets[TORCH] required canonical wrapper binding + targets[JAX] optional binding or explicit TargetGap + parameter_map semantic name -> target API name + output_map semantic role -> target result name + target_only_parameters e.g. Torch current_stream + api_anchors exact public class/wrapper symbols it owns + kernel_anchors exact @cute.kernel symbols it owns + parity_case required when a JAX binding exists +``` + +The target APIs remain separate and explicit: + +```python +# Canonical existing Torch API, unchanged. +from cudnn import rmsnorm_rht_amax_wrapper_sm100 + +torch_result = rmsnorm_rht_amax_wrapper_sm100(...) + +# Optional JAX-native API. +from cudnn.jax import rmsnorm_rht_amax_sm100 + +jax_result = rmsnorm_rht_amax_sm100(...) +``` + +Implicit array-type dispatch is deliberately rejected. JAX tracers, Torch fake +tensors/proxies, mixed operands, and optional framework installations make it +ambiguous and brittle. A `target=` argument inside the call would also become +part of the traced call surface. Importing `cudnn.jax` is the explicit opt-in and +provides an ordinary stable function to `jax.jit`. + +Semantic parity does not mean exact Python signatures or return containers. +Torch legitimately retains `current_stream`, singleton-padding compatibility, +eager allocation, `TupleDict`, and its explicit class lifecycle. JAX owns its +stream and buffers, consumes abstract values, and returns a standard pytree. The +catalog instead checks shared logical operands, common option names/defaults, +output roles, validation/inference, and registered numerical parity coverage. + +The registry requires both target keys. A genuinely unavailable target must be +represented by `TargetGap(reason, tracking_issue)` rather than by omission. The +normal policy for a newly added operation is stricter: new target gaps do not +merge. `TargetGap` exists to inventory pre-existing debt and to make an +exception reviewable, not to make a Torch-only addition look complete. + +Here, "optional JAX" means the dependency and namespace are opt-in for users; +it does not mean JAX support status is optional for contributors. Every +cataloged operation must say implemented or gap, and the no-growth gate expects +new operations to be implemented for JAX. + +The POC's dependency-free AST audit currently finds 64 public Python API +anchors and 58 physical `@cute.kernel` symbols. RMSNorm owns two API anchors and +one kernel under a complete target contract. The remaining 62 API anchors and +57 kernels are exact-symbol migration baselines. Module-level ownership is not +sufficient: it would let a new kernel in an existing module pass unnoticed. +The operation-level JAX-gap baseline starts empty, so a newly cataloged +operation cannot declare JAX missing without an explicit baseline-policy +exception. + +The current wrapper inventory makes the functional gap concrete: + +| Family | Public Torch wrapper entry points | Available JAX bindings | +| --- | ---: | ---: | +| RMSNorm + RHT + amax | 1 | 1 | +| Dense GEMM fusions | 4 | 0 | +| Grouped GEMM | 9 | 0 | +| Discrete grouped GEMM | 2 | 0 | +| SDPA | 2 | 0 | +| Native sparse attention | 4 | 0 | +| DeepSeek sparse attention stages | 11 | 0 | +| **Total** | **33** | **1** | + +This is an inventory, not yet a promise that every legacy helper wrapper should +remain an independent public operation. Phase 1 must group those anchors into +reviewed semantic family/variant rows before generating the long-term support +matrix. + +CI should enforce all of the following: + +1. Discovered public `APIBase` descendants and wrapper functions equal the + union of registered API anchors and the checked-in migration baseline. +2. Every discovered `@cute.kernel` equals a registered kernel owner or a + reviewed internal/migration exception. +3. On a pull request, each baseline set is a subset of its merge-base version. + This is what makes debt non-growing; an equality test against an editable + file alone cannot enforce that policy. +4. Every bound target resolves to a module-level callable whose declared + semantic parameter/default and output-role mappings still match the catalog. + Torch-only parameters are explicit rather than leaked into JAX. +5. Every complete operation supplies common support-domain and numerical parity + cases. Torch and JAX target lanes also test their target-specific lifecycle. +6. JAX cases cover `eval_shape`, `jit`, custom-call lowering, execution, and + documented transformation behavior. + +The AST audit detects structural additions and mapped-default drift. It cannot +infer that an edit inside an already-owned kernel changed semantics. Shared +parity fixtures and target execution tests are therefore mandatory, not +optional follow-up coverage. The canonical Torch wrapper also needs a separate +compatibility snapshot for its signature, `TupleDict` behavior, singleton-rank +handling, and stream semantics. + +The initial scan covers the Python FE-OSS surface. C++ OSS engine registrations +need a separate discovery adapter feeding the same semantic catalog; scanning +Python alone must not be presented as complete repository-wide coverage. + +### 1. Framework-neutral operator core + +Move compile-affecting logic out of framework bindings. The target shape is an +immutable operator definition with these responsibilities: + +```text +OperatorDefinition + config static algorithm/tile/mask attributes + infer_outputs(input_specs) logical output shapes and dtypes + validate(input_specs, target) rank, dtype, layout, divisibility, target limits + make_launcher(config) CuTe launcher or kernel factory + workspace_specs(...) size, alignment, initialization, lifetime + aliasing legal in-place relationships + transform_policy AD, batching, and sharding behavior +``` + +The core must not allocate tensors, inspect device values, access data pointers, +or choose a framework stream. + +The POC starts this extraction with +[`_rmsnorm_rht_amax_config.py`](../../python/cudnn/_rmsnorm_rht_amax_config.py), +which mirrors the canonical Torch launch rules for JAX without modifying the +Torch implementation. Moving both targets onto one neutral implementation is a +Phase 1 change gated by Torch compatibility snapshots. + +### 2. Torch adapter + +Keep existing class and wrapper APIs stable and first class: + +- Convert Torch tensor metadata to canonical operator specs. +- Allocate outputs and eager workspace with Torch. +- Compile through the existing TVM-FFI path. +- Invoke on the explicit/current Torch stream. + +As operator cores are extracted, `APIBase` remains the canonical Torch binding +rather than becoming a framework-generic abstraction. There is no second +`.torch` wrapper, shared result container, or forced signature migration. This +avoids a risky flag day across all existing APIs. + +### 3. JAX adapter + +Expose an optional JAX-native module-level function: + +```python +from cudnn.jax import rmsnorm_rht_amax_sm100 + +@jax.jit +def f(x, weight): + output, amax = rmsnorm_rht_amax_sm100(x, weight) + return output, amax +``` + +Each wrapper: + +1. Reads only abstract shape/dtype metadata. +2. Resolves static configuration. +3. Infers every public output and hidden workspace. +4. Binds a canonical launcher through `call_cutedsl`. +5. Returns a standard tuple, named tuple, or registered pytree. + +Do not dispatch implicitly by checking whether an argument is a Torch tensor or +a JAX tracer. Namespace selection happens before tracing, and imports inside the +JAX implementation keep optional dependencies off the base/Torch path. + +### 4. Future C++ graph adapter + +The C++ cuDNN graph API should not be wrapped in another `cutlass_call`; it has +no Python CuTe compilation to perform. A future `cudnn.jax.graph` layer should +use public typed `jax.ffi.ffi_call` and a registered C++ XLA FFI handler around +the existing sorted-pointer graph execution seam. + +That handler should use stateful FFI stages for immutable plan metadata and a +thread-safe per-device plan/handle cache, but execution must always use XLA's +stream. This work is independent of the FE-OSS CuTe DSL POC. + +## Tensor metadata and layout + +The current `TensorDesc` combines logical tensor metadata with Torch-specific +storage and device types. The production refactor should separate: + +- canonical scalar type; +- logical shape; +- compact physical layout order; +- logical mode permutation presented to the kernel; +- storage packing, especially FP4; +- assumed alignment and dimension divisibility; +- target/device requirements. + +`cutlass.jax.TensorSpec` models compact tensors. It derives runtime strides from +runtime dimensions and the declared layout order; it does not accept arbitrary +runtime stride vectors. The JAX binding must therefore: + +- constrain supported compact layouts in the custom call; +- use `mode` to reinterpret logical modes without materializing a transpose; +- let XLA insert a layout conversion when an operand does not already satisfy + the requested physical layout; +- reject overlapping, broadcast-stride, and non-compact layouts unless the + kernel is explicitly adapted. + +This differs from Torch, where wrappers can allocate an arbitrary +`empty_strided` result. Layout should be an operator contract, not inferred from +a fictitious JAX stride API. + +FP4 also needs a target-specific storage policy. Existing Torch code sometimes +uses packed `uint8` and doubles the logical innermost dimension. JAX/CUTLASS has +a native FP4 dtype mapping. Canonical operator specs must record logical dtype +and packing separately so one target's physical representation does not leak +into another target's shape inference. + +## Compile, cache, and run lifecycle + +### Trace and lowering + +- Shapes, dtypes, optional operand presence, layouts, algorithm choices, tile + sizes, mask modes, and other compile-affecting configuration are static. +- Data-dependent scalars are JAX operands, not Python values. Values such as + tile sizes or a rarely changed `eps` may remain static and cause a recompile. +- No `.item()`, `.cpu()`, `.tolist()`, `.data_ptr()`, DLPack conversion, Torch + allocation, or stream lookup is allowed while tracing. +- Support checks based on shapes/dtypes/layouts run at trace time. Device-value + validation must either become a device computation or an explicit static + bound. + +### Cache keys + +Only immutable compilation state belongs in a cache key: + +- launcher/kernel identity and source version; +- canonical input/output/workspace specs; +- static operator config; +- target architecture and compiler options; +- relevant CUTLASS/cuDNN/frontend ABI versions. + +Do not cache output buffers, workspace buffers, tensor addresses, CUDA streams, +or mutable scheduler counters. CUTLASS currently keys its in-memory compile cache +partly by Python function identity, so the POC memoizes generated launcher +closures. + +CUTLASS 4.5's outer JAX compile-cache key does not visibly include the target +architecture, and both that cache and the POC launcher cache retain entries +without a bound. Before supporting heterogeneous GPU processes or a large +unbounded set of generated launchers, verify device targeting with CUTLASS and +provide coordinated cache observability and clearing. + +### Runtime + +- XLA supplies inputs, outputs, workspaces, and the CUDA stream. +- All helper kernels and the main kernel enqueue on that stream. +- The call returns without synchronizing the host. +- Concurrent executions of one compiled executable must not share mutable + per-call storage. + +## Workspace and initialization design + +Workspace is the main difference from the simple CUTLASS examples. The preferred +representation is an extra custom-call result that the Python wrapper drops: + +```text +custom-call results = public outputs..., hidden workspace buffers... +user-visible result = public outputs... +``` + +OpenXLA documents unused tuple results as +[temporary custom-call buffers](https://openxla.org/xla/custom_call#tuple_outputs_as_temp_buffers). +This lets XLA account for the memory, reuse it according to liveness, preserve +asynchronous lifetime, and avoid runtime allocation during CUDA-graph capture. + +The POC's [`call_cutedsl`](../../python/cudnn/jax/cutedsl.py) supports these +categories: + +| Buffer requirement | JAX representation | +| --- | --- | +| Fully overwritten output/workspace | Uninitialized custom-call result | +| Must start at zero | `jnp.zeros` operand aliased to the result | +| Must start at a constant such as `-inf` | `jnp.full` operand aliased to the result | +| Public in-place result | Declared input/output alias; donation recommended | +| Forward residual needed by backward | Real JAX output/residual, not workspace | +| Data-dependent workspace size | Unsupported initially; use a static upper bound or later runtime allocator fallback | + +An initializer is an aliased operand because custom-call result buffers are not +initialized. XLA retains functional semantics; `donate_argnums` is needed when a +caller wants guaranteed reuse of a user input rather than copy protection. + +Workspace size must be derivable from abstract input metadata before XLA buffer +assignment. If a future plan can only determine its workspace at runtime, choose +one of: + +1. a deterministic safe upper bound; +2. trace/lowering-time plan selection and a serialized, pointer-free recipe; +3. XLA FFI `ScratchAllocator` as a measured fallback. + +Runtime scratch allocation is less visible to XLA's memory planner and needs +separate CUDA-graph, alignment, and asynchronous-lifetime validation. It should +not be the default. + +For one packed byte workspace, overallocate and align internal slices when +operator-specific subregions require stronger alignment. A zero-byte workspace +must not assume a non-null pointer. + +### Persistent descriptor and scheduler workspace + +Several grouped/discrete GEMM APIs keep TMA descriptors or scheduler counters in +a cached Torch tensor. That storage cannot be shared by concurrent JAX calls. +Initially, allocate it per invocation and run descriptor/counter initialization +as helper kernels on XLA's stream before the main kernel. Persistent device state +is only worth adding after profiling proves this cost material, and it would need +a thread-safe stateful runtime design. + +## Operator-specific audit + +| Family | Main JAX work | +| --- | --- | +| RMSNorm + RHT + amax | Static output inference, BF16 validation, per-CTA amax shape. POC complete. | +| Dense GEMM fusions | Map M/N-major compact layouts, FP4/FP8 storage, multiple auxiliary outputs, and initialized amax buffers. | +| SDPA forward | Make max sequence bounds static; remove `.item()` inference; model optional varlen operands and output layouts. | +| SDPA backward | Per-call zeroed workspace, multi-kernel ordering, residual contract, and `custom_vjp`. | +| Grouped GEMM | Per-call descriptor/scheduler workspace, helper launches, counters, and metadata validation without host reads. | +| Discrete grouped GEMM | Current device-address tensors are not traceable. Pass expert buffers as real operands and build pointer tables in device workspace, or defer this mode. | +| DSA/NSA orchestration | Remove host value inspection, define fixed-capacity outputs for compaction, and audit every multi-kernel temporary. | +| Distributed kernels | Replace `torch.distributed` assumptions with explicit JAX collectives/partitioning; do not hide cross-replica state in a CuTe call. | + +Optional tensors alter call arity and are therefore static graph variants. +Data-dependent compacted output lengths require a fixed-capacity result plus a +valid-count output. + +Random operations must accept explicit JAX key/seed/counter operands and return +advanced state where needed. They must not consume hidden global RNG state. + +## JAX transformations + +### Autodiff + +`cutlass_call` explicitly rejects JVP/transpose rules. Inference-only wrappers +should expose that limitation clearly. Operators with forward and backward +kernels should use `jax.custom_vjp`: + +- forward returns the primal outputs plus explicit residual arrays; +- backward consumes residuals and cotangents and invokes a separate CuTe call; +- reserve/stat tensors needed later are residuals, not scratch; +- higher-order and forward-mode differentiation require separate decisions. + +### Batching + +`cutlass_call` explicitly rejects `vmap`, even though the lower-level public FFI +has generic batching modes. Add `custom_vmap` only when an operator defines how +the new batch axis changes shapes, layouts, workspace, and launch configuration. +Do not silently use sequential batching in the FE wrapper. + +### Sharding + +The initial contract is `shard_map`: the FE wrapper sees device-local shapes and +launches one local kernel. Automatic partitioning needs an operator-specific +`custom_partitioning` rule that states: + +- which axes may be sharded without communication; +- whether output and workspace shapes are local or global; +- which layouts remain valid after partitioning; +- what collectives or resharding are required. + +Do not claim transparent multi-device support solely because a custom call can +be replicated. + +### Export and dynamic shapes + +CUTLASS demonstrates `jax.export` with symbolic dimensions and divisibility +hints, but each FE kernel must be audited. The POC intentionally requires +concrete `M` and `N` and uses static tensors. + +Exported programs also require the matching CUTLASS runtime and custom-call +registration in the consuming process. They are not runtime-independent +StableHLO artifacts. + +## CUDA graphs and effects + +Normal FE operations are pure: they read declared operands and write declared +results/aliases. Keep side effects false so XLA may optimize them normally. +Global state, logging, or hidden RNG mutation would violate this contract. + +`cutlass_call` enables CUDA graphs by default. An operator should opt out when it +performs unsupported first-use initialization, runtime allocation, host +synchronization, or dynamic control. Tests must warm compilation and module +loading before measuring capture/replay. + +## Public API and packaging + +Recommended public shape: + +```python +# Existing canonical Torch API and return type remain unchanged. +from cudnn import rmsnorm_rht_amax_wrapper_sm100 + +torch_result = rmsnorm_rht_amax_wrapper_sm100( + x_tensor=x_torch, + w_tensor=weight_torch, + eps=1e-5, + num_threads=128, + rows_per_cta=2, + current_stream=None, +) + +# JAX is a separate optional install and namespace. +import jax +from cudnn.jax import rmsnorm_rht_amax_sm100 + +eps = 1e-5 +num_threads = 128 +rows_per_cta = 2 + +@jax.jit +def run(x, weight): + return rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + +jax_result = run(x_jax, weight_jax) +``` + +- Torch remains the default, first-class API. Its class lifecycle, wrapper + signature, `TupleDict`, singleton-padding behavior, and stream controls are + compatibility contracts. +- JAX is explicitly selected by importing `cudnn.jax`; installing the base or + Torch extras does not require JAX. +- JAX operations use functional inputs/outputs and standard named tuples or + registered pytrees. They do not emulate Torch output buffers or streams. +- Common option names and defaults should be retained when meaningful, but + exact target signatures and containers are not required. +- Compile-affecting JAX keyword arguments are Python-static: close them over as + above or name them in `jax.jit(static_argnames=...)`. Dynamic tensor/scalar + operands remain explicit inputs. +- Clear errors for unsupported transforms and shape modes. +- A flat array ABI at the generic CuTe call seam. Operator-level pytrees may be + flattened and reconstructed by their wrappers, but nested inputs do not cross + the low-level adapter. + +The internal catalog is for ownership, support reporting, and CI; it is not a +replacement public dispatch facade. + +The POC adds a `jax` optional dependency group. It intentionally does +not depend on Torch. Production CI needs a JAX lane whose test bootstrap also +does not unconditionally import Torch. + +JAX follows a fast-moving pre-1.0 compatibility policy. The current official +CuTe DSL + JAX guide lists JAX 0.8.1 as its minimum, while CUTLASS 4.5's module +still declares a looser 0.5.0 minimum. The POC follows the current guide, +bounds the experimental range below JAX 0.10, and retains the repository's +CUTLASS 4.5.0 pin. JAX 0.8.1 and newer require Python 3.11, so this extra has a +narrower Python requirement than the base package. The repository's current +test requirements also pin NumPy below 2.0, while JAX 0.8.1 requires NumPy 2; +the JAX lane therefore needs a separate dependency bootstrap. Pin and test a +supported JAX/CUTLASS/Python matrix before promoting the integration. In +particular, CUTLASS 4.5 selects its original FFI implementation below JAX 0.9.1 +and its stateful FFI implementation at JAX 0.9.1 or newer, so CI must cover both +sides of that boundary. + +## Proof of concept + +### Implemented pieces + +- [`cudnn.frontend`](../../python/cudnn/frontend) + - is an internal dependency-light support and ownership catalog, not a public + execution facade; + - requires a concrete canonical Torch wrapper plus a JAX binding or explicit + tracked gap; + - records semantic parameter/output mappings and intentional target-only + parameters such as `current_stream`; + - requires a parity-case identifier whenever JAX is available; + - records exact public API and physical kernel ownership anchors; + - includes a dependency-free audit of all current Python FE-OSS APIs and + `@cute.kernel` definitions. +- [`cudnn.jax.cutedsl`](../../python/cudnn/jax/cutedsl.py) + - remains an internal POC seam rather than a top-level public export until a + real workspace-using operator validates the contract; + - translates output/workspace metadata to `cutlass_call`; + - appends hidden workspaces and drops them from public results; + - supports uninitialized, zeroed, and constant-filled buffers; + - reconstructs canonical launcher order when results alias inputs; + - requires identical input/output metadata for an alias and rejects nested + low-level input pytrees; + - forwards compact layout/mode/divisibility/alignment metadata; + - memoizes launcher adapters for stable CUTLASS cache identity. +- [`cudnn.jax.rmsnorm_rht_amax_sm100`](../../python/cudnn/jax/rmsnorm_rht_amax.py) + - validates abstract rank, shape, and BF16 dtype; + - matches the canonical Torch option names/defaults and launch heuristics; + - declares `(M, N)` BF16 and `(M / rows_per_cta,)` FP32 outputs; + - lowers the existing CuTe kernel through `cutlass_call`. +- CPU-only adapter contract tests cover hidden workspace, initialization, + aliases, layouts, and invalid metadata. +- CPU-only parity-contract tests cover canonical Torch ownership, optional JAX + status, mapped-default drift, parity-case presence, API discovery, and exact + kernel ownership. +- A GPU test inspects StableHLO, executes one compiled program twice, and checks + a numerical JAX reference when JAX, CuTe DSL, CUDA, and SM100 hardware are + available. + +### POC limitations + +- Only one operation is exposed. +- Existing APIs and physical kernels remain in explicit migration baselines; + the POC has not yet grouped all legacy anchors into semantic operation rows. +- The local AST test makes baseline edits visible but cannot by itself forbid a + contributor from adding a new exception. Production CI must compare the + baseline sets with the pull request's merge base and reject growth. +- `M` and `N` must be concrete during tracing. +- JAX currently accepts the rank-2/rank-1 core domain only; the canonical Torch + wrapper's singleton-padded compatibility remains Torch-specific. +- No custom gradient, batching, or partitioning rule. +- No local SM100 GPU was available for execution in the development + environment; the GPU test is present but must run in CI. +- The generic workspace path is contract-tested with a fake CUTLASS bridge; the + first real operator does not need device workspace. +- The JAX optional dependency version range needs release-owner approval. +- Launcher and CUTLASS compile caches are currently unbounded. + +## Rollout plan and acceptance gates + +### Phase 0: vertical-slice POC + +Status: implemented in this change. + +- Add the JAX CuTe call adapter. +- Extract one framework-neutral launch config. +- Port RMSNorm + RHT + amax. +- Add the Torch-canonical support catalog and exact-symbol ownership audit. +- Add CPU contract tests and an SM100 numerical test. + +Gate: run the GPU test on SM100, inspect lowered StableHLO for one custom call, +and verify a second invocation reuses the compiled executable. + +### Phase 1: neutral metadata core + +- Replace the flat migration baselines with catalog rows grouped by semantic + family/variant, each carrying API/kernel ownership and gap issue metadata. +- Add merge-base non-growth enforcement and generate a checked support matrix + from the catalog. +- Define canonical dtype plus explicit FP4 packing. +- Define logical shape, compact physical layout, mode mapping, alignment, and + divisibility independent of Torch. +- Split shape/output inference and support validation from allocation. +- Preserve existing Torch constructors, wrappers, result containers, and dtype + arguments as independently tested compatibility contracts. +- Add Torch/JAX cross-target metadata tests. + +Gate: the RMSNorm operator has one validation/inference implementation, both +targets retain numerical behavior, and a new API or `@cute.kernel` fails CI +until it has a complete semantic owner. + +### Phase 2: workspace-free dense kernels + +- Port one dense GEMM fusion with multiple outputs. +- Validate M/N-major layout conversion and FP8. +- Exercise initialized auxiliary outputs and legal donation. +- Define compile-cache observability and limits. + +Gate: JIT correctness, layout HLO inspection, concurrent calls, cache reuse, and +Torch parity on supported shapes. + +### Phase 3: workspace and training + +- Port SDPA forward/backward. +- Allocate per-call hidden workspace with explicit zero/fill semantics. +- Make sequence maxima static arguments or explicit bounds. +- Add `custom_vjp` and residual outputs. +- Validate CUDA graph replay and reentrancy on multiple streams. + +Gate: forward/backward numerical parity, no shared mutable workspace, gradient +tests, concurrent execution, and memory/capture tests. + +### Phase 4: grouped, discrete, and sharded operations + +- Build descriptor and pointer tables from real operands on device. +- Remove host reads of tensor metadata values. +- Add operator-specific `shard_map` examples and partitioning rules. +- Integrate explicit JAX collectives where distributed semantics require them. + +Gate: multi-device correctness, pointer lifetime safety, no host synchronization +in the runtime path, and deterministic workspace bounds. + +### Separate track: C++ graph API + +- Define a versioned deterministic graph/plan descriptor. +- Register a typed GPU XLA FFI target. +- Use a thread-safe per-device state/handle strategy. +- Declare logical outputs and hidden workspace through `jax.ffi.ffi_call`. + +This track should reuse neutral tensor/operator metadata but not the CuTe +compilation adapter. + +## Risks and decisions still needing owner approval + +- Supported JAX/CUTLASS/CUDA version matrix and wheel packaging. +- Stable canonical dtype/packing representation, particularly FP4. +- Whether JAX operation naming and named-result conventions match broader cuDNN + Python API plans. +- Which dynamic-shape/export guarantees are required per operation. +- Whether any operator's workspace is genuinely data-dependent after a safe + plan policy is fixed. +- Performance threshold for reinitializing grouped-GEMM descriptor workspace on + every invocation. +- Per-operator AD, batching, sharding, aliasing, and CUDA-graph policies. + +These choices should be recorded per operator rather than hidden in a generic +adapter default. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 0a7b8774a..aadb46e5f 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -2,6 +2,9 @@ **FE-OSS APIs are experimental and subject to change.** +Design work for tracing FE-OSS CuTe DSL kernels from JAX is documented in +[CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md). + This folder documents the Python FE APIs implemented under `python/cudnn`. For details on currently implemented operations, see: - [GEMM + Amax](gemm_fusions/gemm_amax.md) - [GEMM + SwiGLU](gemm_fusions/gemm_swiglu.md) @@ -34,9 +37,32 @@ pip install nvidia-cudnn-frontend[cutedsl] After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` +The experimental JAX proof of concept uses a separate optional dependency set: + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +This optional dependency set requires Python 3.11 or newer, matching the JAX +CUDA package requirement; the base frontend package retains its broader Python +support. + ## API Usage -Each operation exposes two APIs: +Torch remains the canonical FE-OSS API and exposes the following wrapper and +class styles unchanged. Operations with optional JAX support are exposed in a +separate namespace: + +```python +from cudnn import operation_name_wrapper # canonical Torch API +from cudnn.jax import operation_name # optional JAX API +``` + +JAX is additive: it does not replace or narrow the Torch wrapper, `TupleDict`, +class lifecycle, singleton-padding compatibility, or stream controls. The two +APIs share semantic options and kernel behavior where meaningful, while JAX +uses functional results and an XLA-owned stream. Backend selection is never +inferred from array types. ### 1. High-level wrapper diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index ba875fef8..030cc9ed3 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -9,6 +9,7 @@ This frontend integration exposes the kernel as a standard FE-OSS Python API with: - a class API (`RmsNormRhtAmaxSm100`) - a wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) +- an optional experimental JAX API (`cudnn.jax.rmsnorm_rht_amax_sm100`) - grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse ## Shapes @@ -97,6 +98,47 @@ op.execute( ) ``` +### Optional experimental JAX API + +The existing Torch APIs above remain canonical and unchanged. Install the +JAX-specific optional dependencies to use the functional JAX namespace: + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +The JAX optional dependency set requires Python 3.11 or newer. + +```python +import jax +from cudnn.jax import rmsnorm_rht_amax_sm100 + +eps = 1e-5 +num_threads = 128 +rows_per_cta = 2 + +@jax.jit +def run(x, weight): + return rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + +o, amax = run(x, weight) +``` + +The proof of concept requires concrete `M` and `N` during tracing and does not +yet define autodiff, `vmap`, or automatic partitioning rules. `eps`, +`num_threads`, and `rows_per_cta` are static compilation state; close them over +as above or list them in `jax.jit(static_argnames=...)`. JAX returns a standard +`RmsNormRhtAmaxResult(output, amax)` named tuple; Torch retains its existing +`TupleDict(o_tensor, amax_tensor)` contract. See +[CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md) for the design, +workspace model, rollout plan, and current limitations. + ## Parameters ### Input and output tensors @@ -125,7 +167,7 @@ op.execute( - Threads per CTA. If omitted, the API uses the upstream-tuned table when possible, otherwise a valid fallback search. - `rows_per_cta: Optional[int]` - Rows processed by each CTA. If omitted, the wrapper uses the upstream-style heuristic over `{2, 4, 8}`. -- CUDA stream (`current_stream`) +- Torch-only CUDA stream (`current_stream`); JAX always uses XLA's stream ### Wrapper return values @@ -143,6 +185,7 @@ Tuple unpacking order is `(o_tensor, amax_tensor)`. - `EPT = N / num_threads` must be at least `8` and divisible by `8`. - `M` must be divisible by `rows_per_cta`. - Inputs and output are currently bf16 only. +- The JAX proof of concept requires concrete shapes and compact row-major inputs. - The frontend integration matches the upstream RMSNorm kernel semantics; it does not expose full LayerNorm mean/bias behavior. ## Verification diff --git a/pyproject.toml b/pyproject.toml index ca8243c38..f2f5b1239 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,3 +76,8 @@ cutedsl = [ "apache-tvm-ffi", "torch-c-dlpack-ext", ] +jax = [ + "nvidia-cutlass-dsl[cu13]==4.5.0", + "cuda-python", + "jax[cuda13]>=0.8.1,<0.10", +] diff --git a/python/cudnn/README.md b/python/cudnn/README.md index f574cd667..7854e7af6 100644 --- a/python/cudnn/README.md +++ b/python/cudnn/README.md @@ -14,6 +14,8 @@ A simplified view of package structure: pyproject.toml # Project metadata and dependencies. Optional dependencies for frontend-only APIs are registered here. python/cudnn/ ├── __init__.py # Top-level exports (Graph, graph, jit, wrappers, kernels) +├── frontend/ # Internal target-support and kernel-ownership catalog +├── jax/ # Optional functional JAX APIs ├── graph.py # Low-level graph helpers (graph, jit, graph_cache) ├── wrapper.py # High-level Graph wrapper class ├── datatypes.py # Data type conversions and helpers @@ -26,15 +28,51 @@ test/python/ # Test files └── fe_api/ # Test files for frontend-only APIs ``` -## - ## Adding new frontend-only APIs -To add a new frontend-only API, follow these steps: -1. Create a new directory in the `python/cudnn` directory with the name of the API. -2. Add your kernel implementation and implement the high level API implementation in `api.py`, extending the `APIBase` class in `api_base.py`. -3. Expose the API import in `python/cudnn/__init__.py` and register the folder in `pyproject.toml`. Register any optional dependences if required. -4. Add a sample usage/test file in `test/python/fe_api/`. +The intended unit is a semantic operation variant, which may own one or more +main/helper CuTe kernels. The existing Torch conventions remain canonical; JAX +is an optional additive namespace: + +```python +from cudnn import my_operation_wrapper +torch_result = my_operation_wrapper(torch_inputs, ...) + +from cudnn.jax import my_operation +jax_result = my_operation(jax_inputs, ...) +``` + +To add a new frontend-only API: + +1. Define the semantic operation: logical inputs and outputs, common option + names/defaults, shape/dtype/layout inference, support domain, workspace, + aliasing, and initialization behavior. +2. Add the CuTe implementation. Every `@cute.kernel`, including helper kernels, + must be owned by the semantic operation's exact `kernel_anchors` entry. +3. Preserve or add the canonical Torch class/wrapper API following existing + conventions. Do not replace its `TupleDict`, stream controls, output-buffer + lifecycle, or compatibility behavior to make it resemble JAX. +4. Add the functional JAX adapter using `cutlass.jax.cutlass_call`. It must infer + outputs/workspace from abstract metadata, accept XLA's stream, and avoid + Torch imports or host reads during tracing. +5. Register the internal semantic contract in `python/cudnn/frontend`. The Torch + binding is required; JAX must be a `TargetBinding` or explicit + `TargetGap(reason, tracking_issue)`. Record parameter/output mappings, + target-only arguments, exact `api_anchors`, and exact `kernel_anchors`. +6. Add common support-domain/numerical parity cases plus Torch and JAX lifecycle + tests in `test/python/fe_api/`. JAX coverage should include `eval_shape`, + `jit`, lowering, and execution on supported hardware. +7. Run `test_frontend_target_parity.py`. A new public class/wrapper or physical + kernel without a semantic owner is a failure. Existing baselines are + migration debt and must not grow for a new operation. +8. Keep the Torch exports in `cudnn` and expose JAX only from `cudnn.jax`. + Register JAX dependencies only in the optional extra. + +The catalog is for ownership, support reporting, and CI; it is not a public +backend-dispatch facade. Do not use tensor-type dispatch or a traced `target=` +argument. JAX does not emulate the mutable `APIBase.compile()` / `execute()` +lifecycle. JAX is optional for users to install, but declaring its support +status is mandatory for contributors; the normal policy rejects new JAX gaps. **Currently implemented frontend-only APIs**: - `GEMM + Amax` diff --git a/python/cudnn/_rmsnorm_rht_amax_config.py b/python/cudnn/_rmsnorm_rht_amax_config.py new file mode 100644 index 000000000..9e1864fe3 --- /dev/null +++ b/python/cudnn/_rmsnorm_rht_amax_config.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral launch config for RMSNorm + RHT + amax.""" + +from typing import Optional, Tuple + +DEFAULT_NUM_THREADS_BY_N = { + 2048: 128, + 4096: 256, + 7168: 128, + 8192: 512, + 16384: 1024, + 32768: 512, +} +RPC_CANDIDATES = (2, 4, 8) +TARGET_MIN_CTAS = 148 + + +def best_num_threads(n: int) -> Optional[int]: + for num_threads in (1024, 512, 256, 128, 64): + if n % num_threads != 0: + continue + ept = n // num_threads + if ept >= 8 and ept % 8 == 0: + return num_threads + return None + + +def pick_rows_per_cta(m: int) -> int: + for rows_per_cta in reversed(RPC_CANDIDATES): + if m % rows_per_cta != 0: + continue + num_ctas = m // rows_per_cta + if num_ctas >= TARGET_MIN_CTAS: + return rows_per_cta + return RPC_CANDIDATES[0] + + +def resolve_launch_config( + m: int, + n: int, + *, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> Tuple[int, int]: + """Validate dimensions and resolve target-independent launch parameters.""" + + if m <= 0: + raise ValueError(f"M must be positive, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if n % 16 != 0: + raise ValueError(f"N must be divisible by 16 for the Hadamard block size, got {n}") + + resolved_num_threads = num_threads + if resolved_num_threads is None: + resolved_num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) + if resolved_num_threads is None: + raise ValueError(f"No valid num_threads found for N={n}") + if resolved_num_threads <= 0: + raise ValueError(f"num_threads must be positive, got {resolved_num_threads}") + if resolved_num_threads % 32 != 0: + raise ValueError(f"num_threads must be warp-aligned, got {resolved_num_threads}") + if resolved_num_threads > 1024: + raise ValueError("num_threads must not exceed the CUDA block size limit, " f"got {resolved_num_threads}") + + resolved_rows_per_cta = rows_per_cta + if resolved_rows_per_cta is None: + resolved_rows_per_cta = pick_rows_per_cta(m) + if resolved_rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {resolved_rows_per_cta}") + if m % resolved_rows_per_cta != 0: + raise ValueError("M must be divisible by rows_per_cta, " f"got M={m}, rows_per_cta={resolved_rows_per_cta}") + if n % resolved_num_threads != 0: + raise ValueError(f"N={n} must be divisible by num_threads={resolved_num_threads}") + + ept = n // resolved_num_threads + if ept < 8 or ept % 8 != 0: + raise ValueError(f"EPT={ept} must be >= 8 and divisible by 8") + + return resolved_num_threads, resolved_rows_per_cta + + +__all__ = [ + "DEFAULT_NUM_THREADS_BY_N", + "RPC_CANDIDATES", + "TARGET_MIN_CTAS", + "best_num_threads", + "pick_rows_per_cta", + "resolve_launch_config", +] diff --git a/python/cudnn/frontend/__init__.py b/python/cudnn/frontend/__init__.py new file mode 100644 index 000000000..4c5f893d6 --- /dev/null +++ b/python/cudnn/frontend/__init__.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Internal support catalog for Torch-first frontend-only operations.""" + +from ._legacy_gaps import ( + PREEXISTING_JAX_GAP_ANCHORS, + PREEXISTING_KERNEL_OWNERSHIP_GAPS, + PREEXISTING_REGISTERED_JAX_GAP_IDS, +) +from ._registry import FRONTEND_OPERATION_REGISTRY, FrontendTarget +from .rmsnorm_rht_amax import ( + _rmsnorm_rht_amax_sm100_contract as _rmsnorm_rht_amax_sm100_contract, +) + + +def registered_operations(): + """Return registered Torch-canonical operation specifications.""" + + return FRONTEND_OPERATION_REGISTRY.operations() + + +def known_jax_gaps(): + """Return the checked-in migration baseline of Torch-only API anchors.""" + + return tuple(sorted(PREEXISTING_JAX_GAP_ANCHORS)) + + +def known_kernel_ownership_gaps(): + """Return physical CuTe kernels not yet mapped to semantic operations.""" + + return tuple(sorted(PREEXISTING_KERNEL_OWNERSHIP_GAPS)) + + +def known_registered_jax_gap_ids(): + """Return operation-level JAX gaps grandfathered into the catalog.""" + + return tuple(sorted(PREEXISTING_REGISTERED_JAX_GAP_IDS)) + + +__all__ = [ + "FrontendTarget", + "known_jax_gaps", + "known_kernel_ownership_gaps", + "known_registered_jax_gap_ids", + "registered_operations", +] diff --git a/python/cudnn/frontend/_legacy_gaps.py b/python/cudnn/frontend/_legacy_gaps.py new file mode 100644 index 000000000..58266fbc0 --- /dev/null +++ b/python/cudnn/frontend/_legacy_gaps.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Checked-in baseline of FE-OSS API anchors not yet migrated to JAX. + +This list is intentionally explicit. API discovery tests reject a newly added +public ``APIBase`` class or wrapper unless it is owned by a registered semantic +operation with an explicit JAX status or added here as migration debt. The +baseline should shrink and should not be treated as the normal path for new +operations. + +Kernel ownership is tracked separately because one semantic operation may use +several main/helper kernels. A new physical ``@cute.kernel`` must be attached +to a registered operation; the second baseline only records migration debt +that existed when target-parity enforcement was introduced. +""" + +PREEXISTING_JAX_GAP_ANCHORS = frozenset( + { + "cudnn.deepseek_sparse_attention.indexer_backward.api:DenseIndexerBackward", + "cudnn.deepseek_sparse_attention.indexer_backward.api:IndexerBackward", + "cudnn.deepseek_sparse_attention.indexer_backward.api:dense_indexer_backward_wrapper", + "cudnn.deepseek_sparse_attention.indexer_backward.api:indexer_backward_wrapper", + "cudnn.deepseek_sparse_attention.indexer_forward.api:IndexerForward", + "cudnn.deepseek_sparse_attention.indexer_forward.api:indexer_forward_wrapper", + "cudnn.deepseek_sparse_attention.indexer_top_k.api:IndexerTopK", + "cudnn.deepseek_sparse_attention.indexer_top_k.api:compactify_wrapper", + "cudnn.deepseek_sparse_attention.indexer_top_k.api:indexer_top_k_wrapper", + "cudnn.deepseek_sparse_attention.indexer_top_k.api:local_to_global_wrapper", + "cudnn.deepseek_sparse_attention.score_recompute.api:DenseAttnScoreRecompute", + "cudnn.deepseek_sparse_attention.score_recompute.api:DenseIndexerScoreRecompute", + "cudnn.deepseek_sparse_attention.score_recompute.api:SparseAttnScoreRecompute", + "cudnn.deepseek_sparse_attention.score_recompute.api:SparseIndexerScoreRecompute", + "cudnn.deepseek_sparse_attention.score_recompute.api:dense_attn_score_recompute_wrapper", + "cudnn.deepseek_sparse_attention.score_recompute.api:dense_indexer_score_recompute_wrapper", + "cudnn.deepseek_sparse_attention.score_recompute.api:sparse_attn_score_recompute_wrapper", + "cudnn.deepseek_sparse_attention.score_recompute.api:sparse_indexer_score_recompute_wrapper", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.api:SparseAttentionBackward", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.api:sparse_attention_backward_wrapper", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.api:DiscreteGroupedGemmDswigluSm100", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.api:discrete_grouped_gemm_dswiglu_wrapper_sm100", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_swiglu.api:DiscreteGroupedGemmSwigluSm100", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_swiglu.api:discrete_grouped_gemm_swiglu_wrapper_sm100", + "cudnn.gemm_amax.api:GemmAmaxSm100", + "cudnn.gemm_amax.api:gemm_amax_wrapper_sm100", + "cudnn.gemm_dsrelu.api:GemmDsreluSm100", + "cudnn.gemm_dsrelu.api:gemm_dsrelu_wrapper_sm100", + "cudnn.gemm_srelu.api:GemmSreluSm100", + "cudnn.gemm_srelu.api:gemm_srelu_wrapper_sm100", + "cudnn.gemm_swiglu.api:GemmSwigluSm100", + "cudnn.gemm_swiglu.api:gemm_swiglu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_dglu.api:GroupedGemmDgluSm100", + "cudnn.grouped_gemm.grouped_gemm_dglu.api:grouped_gemm_dglu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_dsrelu.api:GroupedGemmDsreluSm100", + "cudnn.grouped_gemm.grouped_gemm_dsrelu.api:grouped_gemm_dsrelu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_dswiglu.api:GroupedGemmDswigluSm100", + "cudnn.grouped_gemm.grouped_gemm_dswiglu.api:grouped_gemm_dswiglu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_glu.api:GroupedGemmGluSm100", + "cudnn.grouped_gemm.grouped_gemm_glu.api:grouped_gemm_glu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_glu_hadamard.api:GroupedGemmGluHadamardSm100", + "cudnn.grouped_gemm.grouped_gemm_glu_hadamard.api:grouped_gemm_glu_hadamard_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_quant.api:GroupedGemmQuantSm100", + "cudnn.grouped_gemm.grouped_gemm_quant.api:grouped_gemm_quant_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_srelu.api:GroupedGemmSreluSm100", + "cudnn.grouped_gemm.grouped_gemm_srelu.api:grouped_gemm_srelu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_swiglu.api:GroupedGemmSwigluSm100", + "cudnn.grouped_gemm.grouped_gemm_swiglu.api:grouped_gemm_swiglu_wrapper_sm100", + "cudnn.grouped_gemm.grouped_gemm_wgrad.api:GroupedGemmWgradSm100", + "cudnn.grouped_gemm.grouped_gemm_wgrad.api:grouped_gemm_wgrad_wrapper_sm100", + "cudnn.native_sparse_attention.compression.api:CompressionAttention", + "cudnn.native_sparse_attention.compression.api:compression_attention_wrapper", + "cudnn.native_sparse_attention.selection.api:SelectionAttention", + "cudnn.native_sparse_attention.selection.api:selection_attention_wrapper", + "cudnn.native_sparse_attention.sliding_window_attention.api:SlidingWindowAttention", + "cudnn.native_sparse_attention.sliding_window_attention.api:sliding_window_attention_wrapper", + "cudnn.native_sparse_attention.top_k.api:TopKReduction", + "cudnn.native_sparse_attention.top_k.api:topk_reduction_wrapper", + "cudnn.sdpa.bwd.api:SdpabwdSm100D256", + "cudnn.sdpa.bwd.api:sdpa_bwd_wrapper_sm100_d256", + "cudnn.sdpa.fwd.api:SdpafwdSm100D256", + "cudnn.sdpa.fwd.api:sdpa_fwd_wrapper_sm100_d256", + } +) + + +# No semantic operation rows had been migrated with a declared JAX gap when +# this catalog was introduced. CI must compare this set with the merge base so +# it can shrink but cannot be used to admit a new Torch-only operation. +PREEXISTING_REGISTERED_JAX_GAP_IDS = frozenset() + + +PREEXISTING_KERNEL_OWNERSHIP_GAPS = frozenset( + { + "cudnn.deepseek_sparse_attention.indexer_backward.dense_indexer_backward_sm100:DenseIndexerBackward2QGemmSm100.kernel_gemm_dense_2q", + "cudnn.deepseek_sparse_attention.indexer_backward.dense_indexer_backward_sm100:ScoreGradDense.kernel_score_grad", + "cudnn.deepseek_sparse_attention.indexer_backward.dense_indexer_backward_sm90:ScoreGradDenseSm90.kernel_score_grad", + "cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm100:IndexerBackwardSm100.kernel_gemm", + "cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm100:ScoreGradSm100.kernel_score_grad", + "cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm90:IndexerBackwardSm90.kernel", + "cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm90:ScoreGradSm90.kernel_score_grad", + "cudnn.deepseek_sparse_attention.indexer_forward.indexer_fwd_sm100:IndexerForwardSm100.kernel", + "cudnn.deepseek_sparse_attention.indexer_forward.indexer_fwd_sm90:IndexerForwardSm90.kernel", + "cudnn.deepseek_sparse_attention.indexer_top_k.block_scan:block_prefix_sum", + "cudnn.deepseek_sparse_attention.indexer_top_k.compactify:CompactifyKernel.kernel", + "cudnn.deepseek_sparse_attention.indexer_top_k.indexer_top_k_decode_varlen:ComputeDynamicCTAOffsets.compute_offsets_kernel", + "cudnn.deepseek_sparse_attention.indexer_top_k.indexer_top_k_decode_varlen:IndexerTopKKernelVarlenDecode.indexer_topk_kernel", + "cudnn.deepseek_sparse_attention.indexer_top_k.local_to_global_dsl:LocalToGlobalTopK.kernel", + "cudnn.deepseek_sparse_attention.score_recompute.dense_score_recompute_sm100:DenseScoreRecomputeSm100.kernel", + "cudnn.deepseek_sparse_attention.score_recompute.dense_score_recompute_sm90:DenseScoreRecomputeSm90.kernel", + "cudnn.deepseek_sparse_attention.score_recompute.sparse_score_recompute_sm100:SparseScoreRecomputeSm100.kernel", + "cudnn.deepseek_sparse_attention.score_recompute.sparse_score_recompute_sm90:SparseScoreRecomputeSm90.kernel", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100:FlashAttentionDSABackwardSm100.bwd", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100:FlashAttentionDSABackwardSm100.convert", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100:FlashAttentionDSABackwardSm100.sum_OdO", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100:FlashAttentionDSABackwardSm100.sum_dSink", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm90:FlashAttentionDSABackwardSm90.kernel", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm90:_FlashAttentionDSABackwardPostprocessSm90.kernel", + "cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm90:_FlashAttentionDSABackwardPreprocessSm90.kernel", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.discrete_B_blockscaled_grouped_gemm_dglu_dbias:BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel.desc_init_kernel_device_ptrs", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.discrete_B_blockscaled_grouped_gemm_dglu_dbias:BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel.kernel", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_swiglu.discrete_B_blockscaled_grouped_gemm_glu_bias:BlockScaledDiscreteWeightGroupedGemmBiasKernel.desc_init_kernel_device_ptrs", + "cudnn.discrete_grouped_gemm.discrete_grouped_gemm_swiglu.discrete_B_blockscaled_grouped_gemm_glu_bias:BlockScaledDiscreteWeightGroupedGemmBiasKernel.kernel", + "cudnn.gemm_amax.dense_blockscaled_gemm_persistent_amax:Sm100BlockScaledPersistentDenseGemmKernel.kernel", + "cudnn.gemm_dsrelu.dense_blockscaled_gemm_persistent_dsrelu_quant:Sm100BlockScaledPersistentDenseGemmKernel.kernel", + "cudnn.gemm_srelu.dense_blockscaled_gemm_persistent_srelu_quant:Sm100BlockScaledPersistentDenseGemmKernel.kernel", + "cudnn.gemm_swiglu.dense_blockscaled_gemm_persistent_swiglu_interleaved_quant:Sm100BlockScaledPersistentDenseGemmKernel.kernel", + "cudnn.gemm_swiglu.dense_gemm_persistent_swiglu:PersistentDenseGemmKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_dglu.moe_blockscaled_grouped_gemm_dglu_dbias:BlockScaledMoEGroupedGemmDgluDbiasKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_dglu.moe_blockscaled_grouped_gemm_dglu_dbias:BlockScaledMoEGroupedGemmDgluDbiasKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_dsrelu.moe_blockscaled_grouped_gemm_dsrelu_quant:BlockScaledMoEGroupedGemmQuantBwdKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_dsrelu.moe_blockscaled_grouped_gemm_dsrelu_quant:BlockScaledMoEGroupedGemmQuantBwdKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_dswiglu.grouped_gemm_dswiglu_quant:BlockScaledContiguousGroupedGemmKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_glu.moe_blockscaled_grouped_gemm_glu_bias:BlockScaledMoEGroupedGemmGluBiasKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_glu.moe_blockscaled_grouped_gemm_glu_bias:BlockScaledMoEGroupedGemmGluBiasKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_glu_hadamard.moe_blockscaled_grouped_gemm_glu_hadamard:BlockScaledMoEGroupedGemmGluHadamardKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_glu_hadamard.moe_blockscaled_grouped_gemm_glu_hadamard:BlockScaledMoEGroupedGemmGluHadamardKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_quant.grouped_gemm_quant:BlockScaledMoEGroupedGemmQuantKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_quant.grouped_gemm_quant:BlockScaledMoEGroupedGemmQuantKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_srelu.moe_blockscaled_grouped_gemm_srelu_quant:BlockScaledMoEGroupedGemmQuantKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_srelu.moe_blockscaled_grouped_gemm_srelu_quant:BlockScaledMoEGroupedGemmQuantKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_swiglu.grouped_gemm_swiglu_quant:BlockScaledContiguousGroupedGemmKernel.kernel", + "cudnn.grouped_gemm.grouped_gemm_wgrad.moe_blockscaled_grouped_gemm_wgrad:BlockScaledMoEGroupedGemmWgradKernel.helper_kernel", + "cudnn.grouped_gemm.grouped_gemm_wgrad.moe_blockscaled_grouped_gemm_wgrad:BlockScaledMoEGroupedGemmWgradKernel.kernel", + "cudnn.native_sparse_attention.compression.fmha:BlackwellFusedMultiHeadAttentionForward.kernel", + "cudnn.native_sparse_attention.selection.NSA_select_attn_fwd_hmma:HopperSelectAttentionFwd.kernel", + "cudnn.native_sparse_attention.top_k.nsa_top_k_reduction_fwd:FineGrainedReductionQK.kernel", + "cudnn.sdpa.bwd.fmha_backward_sm100_2kernel:BlackwellFusedMultiHeadAttentionBackward.sum_OdO", + "cudnn.sdpa.bwd.fmha_dkdv_d256_sm100:BlackwellFusedAttentionDKDVKernel.dkdv_bwd", + "cudnn.sdpa.bwd.fmha_dq_d256_sm100:BlackwellFusedAttentionDQKernel.kernel", + "cudnn.sdpa.fwd.fmha_forward_sm100_d256:BlackwellFusedMultiHeadAttentionForward.kernel", + } +) + + +__all__ = [ + "PREEXISTING_JAX_GAP_ANCHORS", + "PREEXISTING_KERNEL_OWNERSHIP_GAPS", + "PREEXISTING_REGISTERED_JAX_GAP_IDS", +] diff --git a/python/cudnn/frontend/_registry.py b/python/cudnn/frontend/_registry.py new file mode 100644 index 000000000..2d2cc8897 --- /dev/null +++ b/python/cudnn/frontend/_registry.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-light catalog for Torch-first FE-OSS operation support.""" + +from __future__ import annotations + +import importlib +import inspect +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType +from typing import Any, Callable, ClassVar, Mapping, Optional, Sequence, Tuple, Union + + +class FrontendTarget(str, Enum): + """Framework targets recorded by the FE-OSS support catalog.""" + + TORCH = "torch" + JAX = "jax" + + @classmethod + def normalize(cls, value: Union["FrontendTarget", str]) -> "FrontendTarget": + if isinstance(value, cls): + return value + if not isinstance(value, str): + raise TypeError("target must be a Python-static 'torch' or 'jax' value, " f"got {type(value).__name__}") + try: + return cls(value.lower()) + except ValueError as exc: + choices = ", ".join(target.value for target in cls) + raise ValueError(f"Unknown frontend target {value!r}; expected {choices}") from exc + + +@dataclass(frozen=True) +class TargetBinding: + """Lazy target symbol plus its mapping to the semantic operation contract. + + ``parameter_map`` and ``output_map`` use semantic names as keys and the + target API's concrete names as values. ``target_only_parameters`` records + compatibility controls such as a Torch stream that intentionally have no + JAX equivalent. + """ + + module: str + symbol: str + parameter_map: Mapping[str, str] + output_map: Mapping[str, str] + target_only_parameters: Tuple[str, ...] = () + + def __post_init__(self) -> None: + if not self.module or not self.symbol: + raise ValueError("TargetBinding module and symbol must not be empty") + + parameter_map = dict(self.parameter_map) + output_map = dict(self.output_map) + target_only_parameters = tuple(self.target_only_parameters) + for mapping_name, mapping in ( + ("parameter_map", parameter_map), + ("output_map", output_map), + ): + if not mapping or not all(isinstance(key, str) and key and isinstance(value, str) and value for key, value in mapping.items()): + raise ValueError(f"TargetBinding {mapping_name} must contain non-empty names") + if len(set(mapping.values())) != len(mapping): + raise ValueError(f"TargetBinding {mapping_name} target names must be unique") + if len(set(target_only_parameters)) != len(target_only_parameters) or any(not isinstance(name, str) or not name for name in target_only_parameters): + raise ValueError("TargetBinding target_only_parameters must be unique non-empty names") + if set(parameter_map.values()) & set(target_only_parameters): + raise ValueError("Mapped and target-only parameter names must not overlap") + + object.__setattr__(self, "parameter_map", MappingProxyType(parameter_map)) + object.__setattr__(self, "output_map", MappingProxyType(output_map)) + object.__setattr__(self, "target_only_parameters", target_only_parameters) + + @property + def qualified_name(self) -> str: + return f"{self.module}:{self.symbol}" + + def resolve(self) -> Callable[..., Any]: + module = importlib.import_module(self.module) + try: + implementation = getattr(module, self.symbol) + except AttributeError as exc: + raise ImportError(f"Target binding {self.qualified_name} does not exist") from exc + if not callable(implementation): + raise TypeError(f"Target binding {self.qualified_name} is not callable") + return implementation + + +@dataclass(frozen=True) +class TargetGap: + """Explicit, reviewable declaration of a missing optional target.""" + + reason: str + tracking_issue: str + + def __post_init__(self) -> None: + if not self.reason or not self.tracking_issue: + raise ValueError("A target gap requires a reason and tracking issue") + + +TargetStatus = Union[TargetBinding, TargetGap] + + +@dataclass(frozen=True) +class FrontendOperationSpec: + """One Torch-canonical semantic operation and its optional targets.""" + + name: str + contract_signature: inspect.Signature + targets: Mapping[FrontendTarget, TargetStatus] + api_anchors: Tuple[str, ...] + kernel_anchors: Tuple[str, ...] + output_names: Tuple[str, ...] + parity_case: Optional[str] + required_targets: ClassVar[Tuple[FrontendTarget, ...]] = tuple(FrontendTarget) + canonical_target: ClassVar[FrontendTarget] = FrontendTarget.TORCH + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("Frontend operation name must not be empty") + if not self.api_anchors: + raise ValueError(f"{self.name} must declare at least one API anchor") + if not self.kernel_anchors: + raise ValueError(f"{self.name} must own at least one CuTe kernel") + for anchor_kind, anchors in ( + ("API", self.api_anchors), + ("kernel", self.kernel_anchors), + ): + if len(set(anchors)) != len(anchors): + raise ValueError(f"{self.name} has duplicate {anchor_kind} ownership anchors") + if any(":" not in anchor for anchor in anchors): + raise ValueError(f"{self.name} {anchor_kind} anchors must use " "'module:qualified_name' syntax") + if not self.output_names or len(set(self.output_names)) != len(self.output_names): + raise ValueError(f"{self.name} output names must be non-empty and unique") + if self.parity_case is not None and (not isinstance(self.parity_case, str) or not self.parity_case): + raise ValueError(f"{self.name} parity_case must be a non-empty string") + + unsupported_kinds = { + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + } + if any(parameter.kind in unsupported_kinds for parameter in self.contract_signature.parameters.values()): + raise ValueError(f"{self.name} semantic contract must use explicit " "positional-or-keyword and keyword-only parameters") + + normalized_targets = {FrontendTarget.normalize(target): status for target, status in self.targets.items()} + if len(normalized_targets) != len(self.targets): + raise ValueError(f"{self.name} declares the same target more than once") + missing = set(self.required_targets) - set(normalized_targets) + extra = set(normalized_targets) - set(self.required_targets) + if missing or extra: + raise ValueError( + f"{self.name} target declaration mismatch; missing=" f"{sorted(x.value for x in missing)}, extra=" f"{sorted(x.value for x in extra)}" + ) + if not isinstance(normalized_targets[self.canonical_target], TargetBinding): + raise ValueError(f"{self.name} must have a concrete canonical Torch binding") + if not all(isinstance(status, (TargetBinding, TargetGap)) for status in normalized_targets.values()): + raise TypeError(f"{self.name} target entries must be TargetBinding or TargetGap") + + contract_names = set(self.contract_signature.parameters) + output_names = set(self.output_names) + for target, status in normalized_targets.items(): + if not isinstance(status, TargetBinding): + continue + if set(status.parameter_map) != contract_names: + raise ValueError( + f"{self.name} {target.value} parameter map must cover the " + f"semantic contract exactly; expected={sorted(contract_names)}, " + f"got={sorted(status.parameter_map)}" + ) + if set(status.output_map) != output_names: + raise ValueError( + f"{self.name} {target.value} output map must cover semantic " + f"outputs exactly; expected={sorted(output_names)}, " + f"got={sorted(status.output_map)}" + ) + + if isinstance(normalized_targets[FrontendTarget.JAX], TargetBinding): + if not self.parity_case: + raise ValueError(f"{self.name} has a JAX binding but no registered parity case") + elif self.parity_case is not None: + raise ValueError(f"{self.name} cannot declare a parity case without a JAX binding") + + object.__setattr__(self, "targets", MappingProxyType(normalized_targets)) + object.__setattr__(self, "api_anchors", tuple(self.api_anchors)) + object.__setattr__(self, "kernel_anchors", tuple(self.kernel_anchors)) + object.__setattr__(self, "output_names", tuple(self.output_names)) + + def status(self, target: Union[FrontendTarget, str]) -> TargetStatus: + return self.targets[FrontendTarget.normalize(target)] + + def resolve(self, target: Union[FrontendTarget, str]) -> Callable[..., Any]: + normalized_target = FrontendTarget.normalize(target) + status = self.targets[normalized_target] + if isinstance(status, TargetGap): + raise NotImplementedError(f"{self.name} has no {normalized_target.value} implementation: " f"{status.reason} ({status.tracking_issue})") + + implementation = status.resolve() + self._validate_binding_signature(normalized_target, status, implementation) + return implementation + + def _validate_binding_signature( + self, + target: FrontendTarget, + binding: TargetBinding, + implementation: Callable[..., Any], + ) -> None: + actual_signature = inspect.signature(implementation) + actual_parameters = actual_signature.parameters + expected_names = set(binding.parameter_map.values()) | set(binding.target_only_parameters) + if set(actual_parameters) != expected_names: + raise TypeError(f"{self.name} {target.value} binding parameters drifted; " f"expected {sorted(expected_names)}, got {sorted(actual_parameters)}") + + for semantic_name, target_name in binding.parameter_map.items(): + semantic_parameter = self.contract_signature.parameters[semantic_name] + target_parameter = actual_parameters[target_name] + if semantic_parameter.default != target_parameter.default: + raise TypeError( + f"{self.name} {target.value} default for semantic parameter " + f"{semantic_name!r} drifted; expected " + f"{semantic_parameter.default!r}, got {target_parameter.default!r}" + ) + + +class FrontendOperationRegistry: + """Registry used by API discovery and target-support checks.""" + + def __init__(self) -> None: + self._operations: dict[str, FrontendOperationSpec] = {} + + def register(self, operation: FrontendOperationSpec) -> None: + if operation.name in self._operations: + raise ValueError(f"Duplicate frontend operation {operation.name!r}") + for existing in self._operations.values(): + shared_api_anchors = set(operation.api_anchors) & set(existing.api_anchors) + shared_kernel_anchors = set(operation.kernel_anchors) & set(existing.kernel_anchors) + if shared_api_anchors or shared_kernel_anchors: + raise ValueError( + f"{operation.name} and {existing.name} have overlapping " + f"ownership: api={sorted(shared_api_anchors)}, " + f"kernels={sorted(shared_kernel_anchors)}" + ) + self._operations[operation.name] = operation + + def get(self, name: str) -> FrontendOperationSpec: + try: + return self._operations[name] + except KeyError as exc: + raise KeyError(f"Unknown frontend operation {name!r}") from exc + + def operations(self) -> Tuple[FrontendOperationSpec, ...]: + return tuple(self._operations[name] for name in sorted(self._operations)) + + def audit( + self, + *, + require_jax_complete: bool = False, + resolve_bindings: bool = False, + ) -> Tuple[str, ...]: + """Return unresolved bindings and, optionally, declared JAX gaps.""" + + issues = [] + for operation in self.operations(): + for target in operation.required_targets: + status = operation.status(target) + if isinstance(status, TargetGap): + if require_jax_complete: + issues.append(f"{operation.name}:{target.value}: {status.reason} " f"({status.tracking_issue})") + elif resolve_bindings: + try: + operation.resolve(target) + except Exception as exc: + issues.append(f"{operation.name}:{target.value}: {exc}") + return tuple(issues) + + +FRONTEND_OPERATION_REGISTRY = FrontendOperationRegistry() + + +def frontend_operation( + *, + name: str, + targets: Mapping[Union[FrontendTarget, str], TargetStatus], + api_anchors: Sequence[str], + kernel_anchors: Sequence[str], + output_names: Sequence[str], + parity_case: Optional[str], + registry: Optional[FrontendOperationRegistry] = None, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Register a semantic contract without replacing its Python function.""" + + selected_registry = FRONTEND_OPERATION_REGISTRY if registry is None else registry + + def decorate(contract: Callable[..., Any]) -> Callable[..., Any]: + operation = FrontendOperationSpec( + name=name, + contract_signature=inspect.signature(contract), + targets=targets, + api_anchors=tuple(api_anchors), + kernel_anchors=tuple(kernel_anchors), + output_names=tuple(output_names), + parity_case=parity_case, + ) + selected_registry.register(operation) + return contract + + return decorate + + +__all__ = [ + "FRONTEND_OPERATION_REGISTRY", + "FrontendOperationRegistry", + "FrontendOperationSpec", + "FrontendTarget", + "TargetBinding", + "TargetGap", + "frontend_operation", +] diff --git a/python/cudnn/frontend/rmsnorm_rht_amax.py b/python/cudnn/frontend/rmsnorm_rht_amax.py new file mode 100644 index 000000000..18e12a17f --- /dev/null +++ b/python/cudnn/frontend/rmsnorm_rht_amax.py @@ -0,0 +1,69 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Semantic catalog entry for RMSNorm + RHT + amax.""" + +from __future__ import annotations + +from typing import Any, Optional + +from ._registry import FrontendTarget, TargetBinding, frontend_operation + +_COMMON_PARAMETER_NAMES = { + "x": "x", + "weight": "weight", + "eps": "eps", + "num_threads": "num_threads", + "rows_per_cta": "rows_per_cta", +} +_COMMON_OUTPUT_NAMES = { + "output": "output", + "amax": "amax", +} + + +@frontend_operation( + name="rmsnorm_rht_amax_sm100", + targets={ + FrontendTarget.TORCH: TargetBinding( + module="cudnn.rmsnorm_rht_amax.api", + symbol="rmsnorm_rht_amax_wrapper_sm100", + parameter_map={ + **_COMMON_PARAMETER_NAMES, + "x": "x_tensor", + "weight": "w_tensor", + }, + output_map={ + **_COMMON_OUTPUT_NAMES, + "output": "o_tensor", + "amax": "amax_tensor", + }, + target_only_parameters=("current_stream",), + ), + FrontendTarget.JAX: TargetBinding( + module="cudnn.jax.rmsnorm_rht_amax", + symbol="rmsnorm_rht_amax_sm100", + parameter_map=_COMMON_PARAMETER_NAMES, + output_map=_COMMON_OUTPUT_NAMES, + ), + }, + api_anchors=( + "cudnn.rmsnorm_rht_amax.api:RmsNormRhtAmaxSm100", + "cudnn.rmsnorm_rht_amax.api:rmsnorm_rht_amax_wrapper_sm100", + ), + kernel_anchors=("cudnn.rmsnorm_rht_amax.kernel:RMSNormRHTAmaxKernel.kernel",), + output_names=("output", "amax"), + parity_case="rmsnorm_rht_amax", +) +def _rmsnorm_rht_amax_sm100_contract( + x: Any, + weight: Any, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> None: + """Framework-neutral semantic parameters and defaults for this operation.""" + + +__all__ = [] diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py new file mode 100644 index 000000000..bae8e0423 --- /dev/null +++ b/python/cudnn/jax/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX integration for frontend-only CuTe DSL operations.""" + +from .rmsnorm_rht_amax import RmsNormRhtAmaxResult, rmsnorm_rht_amax_sm100 + +__all__ = ["RmsNormRhtAmaxResult", "rmsnorm_rht_amax_sm100"] diff --git a/python/cudnn/jax/cutedsl.py b/python/cudnn/jax/cutedsl.py new file mode 100644 index 000000000..1c9979195 --- /dev/null +++ b/python/cudnn/jax/cutedsl.py @@ -0,0 +1,484 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX tracing adapter for frontend-only CuTe DSL operations. + +This module is intentionally independent of :mod:`cudnn.api_base`. The +existing API base implements an eager, preallocated-output lifecycle for +PyTorch. JAX needs a functional call whose output buffers are described while +the function is being traced. ``call_cutedsl`` bridges that difference using +``cutlass.jax.cutlass_call``. + +The adapter also models temporary device workspaces as hidden custom-call +results. XLA owns those buffers and can reuse their storage according to the +compiled program's liveness analysis. Buffers that must start at zero or at a +constant value are materialized as JAX inputs and aliased to the corresponding +custom-call results. + +This is a proof of concept. Operator wrappers still need target-neutral shape, +dtype, layout, and support inference before they can use this adapter. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from functools import lru_cache +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple + + +class BufferInitialization(str, Enum): + """Initialization policy for an output or workspace buffer.""" + + UNINITIALIZED = "uninitialized" + ZERO = "zero" + VALUE = "value" + + +@dataclass(frozen=True) +class TensorLayout: + """CuTe/JAX layout metadata for one tensor. + + ``layout`` and ``mode`` follow ``cutlass.jax.TensorSpec`` conventions. In + particular, ``layout`` is CuTe's minor-to-major stride-rank permutation, + while ``mode`` changes the logical mode order seen by the kernel without + requesting a physical transpose. + """ + + layout: Optional[Tuple[int, ...]] = None + mode: Optional[Tuple[int, ...]] = None + static: Optional[bool] = None + ptr_assumed_align: int = 256 + divisibility: Optional[Tuple[Optional[int], ...] | int] = None + + def __post_init__(self) -> None: + if self.layout is not None: + object.__setattr__(self, "layout", tuple(self.layout)) + if self.mode is not None: + object.__setattr__(self, "mode", tuple(self.mode)) + if isinstance(self.divisibility, list): + object.__setattr__(self, "divisibility", tuple(self.divisibility)) + if self.static is not None and not isinstance(self.static, bool): + raise TypeError("static must be bool or None") + if not isinstance(self.ptr_assumed_align, int) or isinstance(self.ptr_assumed_align, bool): + raise TypeError("ptr_assumed_align must be an integer") + if self.ptr_assumed_align <= 0 or (self.ptr_assumed_align & (self.ptr_assumed_align - 1)): + raise ValueError("ptr_assumed_align must be a positive power of two") + divisibility_values = self.divisibility if isinstance(self.divisibility, tuple) else (self.divisibility,) + if any(value is not None and (not isinstance(value, int) or isinstance(value, bool) or value <= 0) for value in divisibility_values): + raise ValueError("divisibility entries must be positive integers or None, " f"got {self.divisibility}") + + def validate_rank(self, rank: int, *, name: str) -> None: + for field_name in ("layout", "mode"): + value = getattr(self, field_name) + if value is None: + continue + if len(value) != rank or tuple(sorted(value)) != tuple(range(rank)): + raise ValueError(f"{name} {field_name} must be a permutation of range({rank}), " f"got {value}") + + if isinstance(self.divisibility, tuple) and len(self.divisibility) != rank: + raise ValueError(f"{name} divisibility must have rank {rank}, " f"got {self.divisibility}") + + +@dataclass(frozen=True) +class BufferSpec: + """Shape, dtype, layout, and initialization for a custom-call result. + + A spec passed through ``outputs`` is returned to the caller. A spec passed + through ``workspaces`` is supplied to the CuTe launcher after the public + outputs and then omitted from the JAX-visible result. + """ + + name: str + shape: Tuple[Any, ...] + dtype: Any + tensor_layout: TensorLayout = field(default_factory=TensorLayout) + initialization: BufferInitialization = BufferInitialization.UNINITIALIZED + fill_value: Any = None + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("BufferSpec.name must not be empty") + if not isinstance(self.tensor_layout, TensorLayout): + raise TypeError("BufferSpec.tensor_layout must be a TensorLayout, " f"got {type(self.tensor_layout).__name__}") + object.__setattr__(self, "shape", tuple(self.shape)) + if not isinstance(self.initialization, BufferInitialization): + object.__setattr__( + self, + "initialization", + BufferInitialization(self.initialization), + ) + self.tensor_layout.validate_rank(len(self.shape), name=self.name) + + if self.initialization is BufferInitialization.VALUE: + if self.fill_value is None: + raise ValueError(f"{self.name} uses VALUE initialization but has no fill_value") + elif self.fill_value is not None: + raise ValueError(f"{self.name} has a fill_value but its initialization policy is " f"{self.initialization.value}") + + +@dataclass(frozen=True) +class _CallPlan: + num_user_inputs: int + all_results: Tuple[BufferSpec, ...] + num_public_results: int + # For every result, either the input index containing its aliased storage or + # None when cutlass_call must provide a fresh output buffer. + result_input_sources: Tuple[Optional[int], ...] + input_output_aliases: Tuple[Tuple[int, int], ...] + initialized_result_indices: Tuple[int, ...] + + @property + def num_total_inputs(self) -> int: + return self.num_user_inputs + len(self.initialized_result_indices) + + +def _build_call_plan( + *, + num_user_inputs: int, + outputs: Sequence[BufferSpec], + workspaces: Sequence[BufferSpec], + input_output_aliases: Optional[Mapping[int, int]], +) -> _CallPlan: + outputs = tuple(outputs) + workspaces = tuple(workspaces) + if not outputs: + raise ValueError("call_cutedsl requires at least one public output") + + all_results = outputs + workspaces + if not all(isinstance(spec, BufferSpec) for spec in all_results): + raise TypeError("outputs and workspaces must contain only BufferSpec values") + names = [spec.name for spec in all_results] + if len(set(names)) != len(names): + raise ValueError(f"Buffer names must be unique, got {names}") + + aliases = dict(input_output_aliases or {}) + if len(set(aliases.values())) != len(aliases): + raise ValueError("Each result may alias at most one input") + + result_sources: list[Optional[int]] = [None] * len(all_results) + for input_idx, output_idx in aliases.items(): + if not isinstance(input_idx, int) or isinstance(input_idx, bool) or not isinstance(output_idx, int) or isinstance(output_idx, bool): + raise TypeError("Alias indices must be integers") + if input_idx < 0 or input_idx >= num_user_inputs: + raise ValueError(f"Invalid aliased input index {input_idx}; there are " f"{num_user_inputs} user inputs") + if output_idx < 0 or output_idx >= len(outputs): + raise ValueError(f"Invalid aliased output index {output_idx}; only public outputs " "may alias user inputs") + if outputs[output_idx].initialization is not BufferInitialization.UNINITIALIZED: + raise ValueError(f"{outputs[output_idx].name} cannot both alias a user input and " "request initialization") + result_sources[output_idx] = input_idx + + initialized_indices = [] + next_input_idx = num_user_inputs + for result_idx, spec in enumerate(all_results): + if spec.initialization is BufferInitialization.UNINITIALIZED: + continue + if result_sources[result_idx] is not None: + raise ValueError(f"{spec.name} has more than one storage source") + result_sources[result_idx] = next_input_idx + aliases[next_input_idx] = result_idx + initialized_indices.append(result_idx) + next_input_idx += 1 + + return _CallPlan( + num_user_inputs=num_user_inputs, + all_results=all_results, + num_public_results=len(outputs), + result_input_sources=tuple(result_sources), + input_output_aliases=tuple(sorted(aliases.items())), + initialized_result_indices=tuple(initialized_indices), + ) + + +def _make_launch_adapter_uncached( + fn: Callable[..., None], + num_user_inputs: int, + num_total_inputs: int, + result_input_sources: Tuple[Optional[int], ...], +) -> Callable[..., None]: + """Reconstruct canonical ``inputs, outputs, workspaces`` launcher order.""" + + def launch_adapter(stream: Any, *args: Any, **kwargs: Any) -> None: + input_args = args[:num_total_inputs] + fresh_results = iter(args[num_total_inputs:]) + canonical_results = [] + for source in result_input_sources: + if source is None: + canonical_results.append(next(fresh_results)) + else: + canonical_results.append(input_args[source]) + + try: + next(fresh_results) + except StopIteration: + pass + else: + raise RuntimeError("CuTe launcher received more result buffers than expected") + + fn( + stream, + *input_args[:num_user_inputs], + *canonical_results, + **kwargs, + ) + + return launch_adapter + + +@lru_cache(maxsize=None) +def _make_launch_adapter_cached( + fn: Callable[..., None], + num_user_inputs: int, + num_total_inputs: int, + result_input_sources: Tuple[Optional[int], ...], +) -> Callable[..., None]: + # Reusing the wrapper identity matters because CUTLASS includes the function + # object in its compilation-cache key. + return _make_launch_adapter_uncached( + fn, + num_user_inputs, + num_total_inputs, + result_input_sources, + ) + + +def _make_launch_adapter(fn: Callable[..., None], plan: _CallPlan) -> Callable[..., None]: + try: + return _make_launch_adapter_cached( + fn, + plan.num_user_inputs, + plan.num_total_inputs, + plan.result_input_sources, + ) + except TypeError: + # Callable objects are normally hashable, including @cute.jit functions. + # Falling back keeps the adapter useful for unusual unhashable callables. + return _make_launch_adapter_uncached( + fn, + plan.num_user_inputs, + plan.num_total_inputs, + plan.result_input_sources, + ) + + +def _to_cutlass_tensor_spec(layout: TensorLayout, cutlass_jax: Any) -> Any: + return cutlass_jax.TensorSpec( + layout=layout.layout, + mode=layout.mode, + static=layout.static, + ptr_assumed_align=layout.ptr_assumed_align, + divisibility=layout.divisibility, + ) + + +def _is_default_tensor_layout(layout: TensorLayout) -> bool: + return layout == TensorLayout() + + +def _make_initialized_buffer(spec: BufferSpec, jax_numpy: Any) -> Any: + if spec.initialization is BufferInitialization.ZERO: + return jax_numpy.zeros(spec.shape, dtype=spec.dtype) + if spec.initialization is BufferInitialization.VALUE: + return jax_numpy.full(spec.shape, spec.fill_value, dtype=spec.dtype) + raise AssertionError(f"Unexpected initialization policy for {spec.name}") + + +def _call_cutedsl_with_modules( + fn: Callable[..., None], + inputs: Sequence[Any], + *, + outputs: Sequence[BufferSpec], + workspaces: Sequence[BufferSpec] = (), + input_layouts: Optional[Sequence[Optional[TensorLayout]]] = None, + input_output_aliases: Optional[Mapping[int, int]] = None, + static_args: Optional[Mapping[str, Any]] = None, + allow_cuda_graph: bool = True, + compile_options: Any = None, + use_static_tensors: bool = False, + jax_module: Any, + jax_numpy_module: Any, + cutlass_jax_module: Any, +) -> Tuple[Any, ...]: + """Dependency-injected implementation used by tests and ``call_cutedsl``.""" + + inputs = tuple(inputs) + static_args = dict(static_args or {}) + reserved_static_args = { + "allow_cuda_graph", + "compile_options", + "fn", + "input_mode", + "input_output_aliases", + "input_spec", + "output_mode", + "output_shape_dtype", + "output_spec", + "use_static_tensors", + } + conflicts = sorted(reserved_static_args.intersection(static_args)) + if conflicts: + raise ValueError("static_args contains names reserved by cutlass_call: " + ", ".join(conflicts)) + plan = _build_call_plan( + num_user_inputs=len(inputs), + outputs=outputs, + workspaces=workspaces, + input_output_aliases=input_output_aliases, + ) + + if input_layouts is None: + normalized_input_layouts: Tuple[Optional[TensorLayout], ...] = (None,) * len(inputs) + else: + normalized_input_layouts = tuple(input_layouts) + if len(normalized_input_layouts) != len(inputs): + raise ValueError(f"Expected {len(inputs)} input layout specs, got " f"{len(normalized_input_layouts)}") + + for input_idx, value in enumerate(inputs): + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError( + "call_cutedsl inputs must be a flat sequence of array-like " + f"values with shape and dtype metadata; input #{input_idx} " + f"is {type(value).__name__}" + ) + + for input_idx, layout in enumerate(normalized_input_layouts): + if layout is None: + continue + if not isinstance(layout, TensorLayout): + raise TypeError(f"Input layout #{input_idx} must be TensorLayout or None, " f"got {type(layout).__name__}") + try: + input_rank = len(inputs[input_idx].shape) + except AttributeError as exc: + raise TypeError(f"Input #{input_idx} must expose shape and dtype metadata") from exc + layout.validate_rank(input_rank, name=f"input #{input_idx}") + + # Fail before lowering when a requested alias cannot describe the same + # logical buffer. XLA may still insert a protective copy unless the caller + # donates a user input, but aliased input/result avals must agree. + for input_idx, output_idx in plan.input_output_aliases: + if input_idx >= plan.num_user_inputs: + continue + input_value = inputs[input_idx] + output_spec = plan.all_results[output_idx] + try: + input_shape = tuple(input_value.shape) + input_dtype = input_value.dtype + except AttributeError as exc: + raise TypeError(f"Aliased input #{input_idx} must expose shape and dtype metadata") from exc + if input_shape != output_spec.shape or input_dtype != output_spec.dtype: + raise ValueError( + f"Aliased input #{input_idx} has shape/dtype " + f"{input_shape}/{input_dtype}, but output " + f"{output_spec.name} requires " + f"{output_spec.shape}/{output_spec.dtype}" + ) + input_layout = normalized_input_layouts[input_idx] + effective_input_layout = input_layout or TensorLayout() + if effective_input_layout != output_spec.tensor_layout: + raise ValueError( + f"Aliased input #{input_idx} and output {output_spec.name} must " + "use identical TensorLayout metadata; CUTLASS compiles the " + "aliased buffer from the input spec" + ) + + initialized_buffers = [] + initialized_layouts = [] + for result_idx in plan.initialized_result_indices: + spec = plan.all_results[result_idx] + initialized_buffers.append(_make_initialized_buffer(spec, jax_numpy_module)) + initialized_layouts.append(spec.tensor_layout) + + def convert_layout(layout: Optional[TensorLayout]) -> Any: + return None if layout is None or _is_default_tensor_layout(layout) else _to_cutlass_tensor_spec(layout, cutlass_jax_module) + + cutlass_input_specs = tuple(convert_layout(x) for x in normalized_input_layouts) + cutlass_input_specs += tuple(convert_layout(x) for x in initialized_layouts) + cutlass_output_specs = tuple(convert_layout(x.tensor_layout) for x in plan.all_results) + result_shape_dtypes = tuple(jax_module.ShapeDtypeStruct(x.shape, x.dtype) for x in plan.all_results) + + launcher = _make_launch_adapter(fn, plan) + call = cutlass_jax_module.cutlass_call( + launcher, + output_shape_dtype=result_shape_dtypes, + input_spec=cutlass_input_specs, + output_spec=cutlass_output_specs, + input_output_aliases=dict(plan.input_output_aliases), + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + **static_args, + ) + all_results = call(*inputs, *initialized_buffers) + + # Passing a tuple to output_shape_dtype asks cutlass_call for multiple + # results, including the one-result case. Keep a defensive normalization + # here so a compatible bridge implementation can still be substituted. + if not isinstance(all_results, (tuple, list)): + all_results = (all_results,) + if len(all_results) != len(plan.all_results): + raise RuntimeError(f"CuTe call returned {len(all_results)} buffers; expected " f"{len(plan.all_results)}") + return tuple(all_results[: plan.num_public_results]) + + +def call_cutedsl( + fn: Callable[..., None], + inputs: Sequence[Any], + *, + outputs: Sequence[BufferSpec], + workspaces: Sequence[BufferSpec] = (), + input_layouts: Optional[Sequence[Optional[TensorLayout]]] = None, + input_output_aliases: Optional[Mapping[int, int]] = None, + static_args: Optional[Mapping[str, Any]] = None, + allow_cuda_graph: bool = True, + compile_options: Any = None, + use_static_tensors: bool = False, +) -> Tuple[Any, ...]: + """Invoke a CuTe DSL launcher as a functional JAX operation. + + The launcher's canonical signature is:: + + fn(stream, *inputs, *outputs, *workspaces, **static_args) -> None + + Call this function while tracing a JAX function (normally inside + ``jax.jit``). Public results are returned as a tuple in ``outputs`` order; + workspace results are hidden. All output and workspace sizes must be + derivable from abstract input metadata, not from runtime tensor values. + + ``input_output_aliases`` uses CUTLASS/JAX's ``{input_index: output_index}`` + convention and only addresses public outputs. Inputs must be a flat + sequence of array-like values; nested pytrees are intentionally outside the + FE ABI. An aliased input and output must use identical ``TensorLayout`` + metadata. The JAX function remains functional: callers should use + ``donate_argnums`` when they want XLA to reuse aliased input storage without + a copy. + """ + + try: + import jax + import jax.numpy as jnp + import cutlass.jax as cutlass_jax + except ImportError as exc: + raise ImportError("cudnn.jax requires JAX and the CuTe DSL JAX integration; install " "the 'jax' optional dependencies") from exc + + return _call_cutedsl_with_modules( + fn, + inputs, + outputs=outputs, + workspaces=workspaces, + input_layouts=input_layouts, + input_output_aliases=input_output_aliases, + static_args=static_args, + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + jax_module=jax, + jax_numpy_module=jnp, + cutlass_jax_module=cutlass_jax, + ) + + +__all__ = [ + "BufferInitialization", + "BufferSpec", + "TensorLayout", + "call_cutedsl", +] diff --git a/python/cudnn/jax/rmsnorm_rht_amax.py b/python/cudnn/jax/rmsnorm_rht_amax.py new file mode 100644 index 000000000..71fa1cdda --- /dev/null +++ b/python/cudnn/jax/rmsnorm_rht_amax.py @@ -0,0 +1,151 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for fused RMSNorm + RHT + per-CTA amax.""" + +from __future__ import annotations + +import operator +from functools import lru_cache +from typing import Any, NamedTuple, Optional + +from .._rmsnorm_rht_amax_config import resolve_launch_config + +from .cutedsl import BufferSpec, TensorLayout, call_cutedsl + + +class RmsNormRhtAmaxResult(NamedTuple): + """Functional JAX outputs for RMSNorm + RHT + amax.""" + + output: Any + amax: Any + + +def _concrete_dim(value: Any, name: str) -> int: + try: + return operator.index(value) + except TypeError as exc: + raise NotImplementedError(f"rmsnorm_rht_amax_sm100 currently requires a concrete {name} " "dimension during tracing") from exc + + +def _static_int(value: Optional[int], name: str) -> Optional[int]: + if value is None: + return None + if isinstance(value, bool): + raise TypeError(f"{name} must be an integer, not bool") + try: + return operator.index(value) + except TypeError as exc: + raise TypeError(f"{name} must be Python-static while tracing; close over the value " "or mark the enclosing jax.jit argument static") from exc + + +def _static_float(value: Any, name: str) -> float: + try: + return float(value) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must be Python-static while tracing; close over the value " "or mark the enclosing jax.jit argument static") from exc + + +@lru_cache(maxsize=None) +def _make_launcher( + n: int, + num_threads: int, + rows_per_cta: int, + eps: float, +): + # Keep optional CuTe DSL imports off the cudnn.jax import path. + from cutlass import Float32 + from ..rmsnorm_rht_amax.kernel import RMSNormRHTAmaxKernel + + kernel = RMSNormRHTAmaxKernel( + n=n, + num_threads=num_threads, + eps=eps, + rows_per_cta=rows_per_cta, + ) + + def launch(stream, x, weight, output, amax): + kernel(x, weight, output, amax, Float32(eps), stream) + + return launch + + +def rmsnorm_rht_amax_sm100( + x: Any, + weight: Any, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> RmsNormRhtAmaxResult: + """Apply fused RMSNorm, 16-wide RHT, and per-CTA amax from JAX. + + This functional API is intended for use inside ``jax.jit``. ``x`` must be + a row-major ``bfloat16`` array of shape ``(M, N)`` and ``weight`` a + ``bfloat16`` array of shape ``(N,)``. ``M`` and ``N`` are concrete in this + proof of concept; shape-polymorphic export is a follow-up. + + ``eps``, ``num_threads``, and ``rows_per_cta`` are compile-time + configuration. Close them over a jitted function or list them in + ``jax.jit(static_argnames=...)``. + + Returns ``(output, amax)`` with shapes ``(M, N)`` and + ``(M // rows_per_cta,)`` respectively. + """ + + try: + import jax.numpy as jnp + except ImportError as exc: + raise ImportError("rmsnorm_rht_amax_sm100 requires JAX; install the 'jax' " "optional dependencies") from exc + + if x.ndim != 2: + raise ValueError(f"x must have rank 2, got shape {x.shape}") + if weight.ndim != 1: + raise ValueError(f"weight must have rank 1, got shape {weight.shape}") + + m = _concrete_dim(x.shape[0], "M") + n = _concrete_dim(x.shape[1], "N") + if tuple(weight.shape) != (n,): + raise ValueError(f"weight shape must match the hidden dimension ({n},), " f"got {weight.shape}") + if x.dtype != jnp.bfloat16 or weight.dtype != jnp.bfloat16: + raise ValueError("x and weight must both have dtype bfloat16, " f"got {x.dtype} and {weight.dtype}") + + num_threads = _static_int(num_threads, "num_threads") + rows_per_cta = _static_int(rows_per_cta, "rows_per_cta") + eps = _static_float(eps, "eps") + resolved_num_threads, resolved_rows_per_cta = resolve_launch_config( + m, + n, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + # JAX/XLA owns physical buffers. These specs constrain ordinary row-major + # storage while providing the divisibility facts used by the CuTe kernel. + x_layout = TensorLayout(divisibility=(resolved_rows_per_cta, 16)) + weight_layout = TensorLayout(divisibility=(16,)) + amax_layout = TensorLayout() + + output, amax = call_cutedsl( + _make_launcher(n, resolved_num_threads, resolved_rows_per_cta, eps), + (x, weight), + outputs=( + BufferSpec( + "output", + (m, n), + jnp.bfloat16, + tensor_layout=x_layout, + ), + BufferSpec( + "amax", + (m // resolved_rows_per_cta,), + jnp.float32, + tensor_layout=amax_layout, + ), + ), + input_layouts=(x_layout, weight_layout), + use_static_tensors=True, + ) + return RmsNormRhtAmaxResult(output=output, amax=amax) + + +__all__ = ["RmsNormRhtAmaxResult", "rmsnorm_rht_amax_sm100"] diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 9da5cae15..09ff4877d 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -1,12 +1,14 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - RmsNormRhtAmaxSm100, - best_num_threads, - pick_rows_per_cta, - rmsnorm_rht_amax_wrapper_sm100, -) +"""Torch API exports for RMSNorm + RHT + amax. + +The exports are lazy so importing the framework-neutral CuTe kernel from the +JAX adapter does not import Torch. +""" + +from importlib import import_module +from typing import Any __all__ = [ "RmsNormRhtAmaxSm100", @@ -14,3 +16,13 @@ "pick_rows_per_cta", "rmsnorm_rht_amax_wrapper_sm100", ] + + +def __getattr__(name: str) -> Any: + if name not in __all__: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = import_module(".api", __name__) + value = getattr(module, name) + globals()[name] = value + return value diff --git a/test/python/fe_api/test_frontend_target_parity.py b/test/python/fe_api/test_frontend_target_parity.py new file mode 100644 index 000000000..87f59de1e --- /dev/null +++ b/test/python/fe_api/test_frontend_target_parity.py @@ -0,0 +1,551 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contract tests for Torch-first FE-OSS target support.""" + +from __future__ import annotations + +import ast +import importlib +from pathlib import Path +import re +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_PYTHON_ROOT = _REPO_ROOT / "python" +_CUDNN_ROOT = _PYTHON_ROOT / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_parity_contract_test" +_REQUIRED_DEFAULT = "" + + +def _load_frontend_package(): + parent = types.ModuleType(_TEST_PACKAGE) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = _TEST_PACKAGE + sys.modules[_TEST_PACKAGE] = parent + return importlib.import_module(f"{_TEST_PACKAGE}.frontend") + + +_FRONTEND = _load_frontend_package() +_REGISTRY_MODULE = importlib.import_module(f"{_TEST_PACKAGE}.frontend._registry") + + +def _binding( + module, + symbol, + *, + parameter_map=None, + output_map=None, + target_only_parameters=(), +): + return _REGISTRY_MODULE.TargetBinding( + module=module, + symbol=symbol, + parameter_map=parameter_map or {"x": "x"}, + output_map=output_map or {"output": "output"}, + target_only_parameters=target_only_parameters, + ) + + +def _base_name(base): + if isinstance(base, ast.Name): + return base.id + if isinstance(base, ast.Attribute): + return base.attr + return None + + +def _decorator_name(decorator): + parts = [] + while isinstance(decorator, ast.Attribute): + parts.append(decorator.attr) + decorator = decorator.value + if isinstance(decorator, ast.Name): + parts.append(decorator.id) + return ".".join(reversed(parts)) + + +def _discover_python_api_anchors(): + anchors = set() + for path in sorted(_CUDNN_ROOT.rglob("api.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + module = ".".join(path.with_suffix("").relative_to(_PYTHON_ROOT).parts) + classes = {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)} + + def is_api_class(name, seen=()): + if name in seen: + return False + for base in classes[name].bases: + parent = _base_name(base) + if parent == "APIBase": + return True + if parent in classes and is_api_class(parent, seen + (name,)): + return True + return False + + for name in classes: + if not name.startswith("_") and is_api_class(name): + anchors.add(f"{module}:{name}") + + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and not node.name.startswith("_") and "wrapper" in node.name: + anchors.add(f"{module}:{node.name}") + return anchors + + +def _discover_cute_kernel_anchors(): + anchors = set() + + class KernelVisitor(ast.NodeVisitor): + def __init__(self, module): + self.module = module + self.scope = [] + + def visit_ClassDef(self, node): + self.scope.append(node.name) + self.generic_visit(node) + self.scope.pop() + + def visit_FunctionDef(self, node): + if any(_decorator_name(decorator) == "cute.kernel" for decorator in node.decorator_list): + qualified_name = ".".join((*self.scope, node.name)) + anchors.add(f"{self.module}:{qualified_name}") + self.scope.append(node.name) + self.generic_visit(node) + self.scope.pop() + + visit_AsyncFunctionDef = visit_FunctionDef + + for path in sorted(_CUDNN_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + module = ".".join(path.with_suffix("").relative_to(_PYTHON_ROOT).parts) + KernelVisitor(module).visit(tree) + return anchors + + +def _ast_function_parameters(path, function_name): + tree = ast.parse(path.read_text(), filename=str(path)) + function = next(node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == function_name) + arguments = function.args + parameters = {} + + positional = (*arguments.posonlyargs, *arguments.args) + positional_defaults = (_REQUIRED_DEFAULT,) * (len(positional) - len(arguments.defaults)) + tuple(ast.dump(default) for default in arguments.defaults) + for argument, default in zip(positional, positional_defaults): + parameters[argument.arg] = ("positional", default) + + for argument, default in zip(arguments.kwonlyargs, arguments.kw_defaults): + parameters[argument.arg] = ( + "keyword_only", + _REQUIRED_DEFAULT if default is None else ast.dump(default), + ) + return parameters + + +class FrontendTargetParityTest(unittest.TestCase): + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def test_registered_operation_keeps_torch_canonical_and_jax_optional(self): + operations = _FRONTEND.registered_operations() + self.assertEqual( + [operation.name for operation in operations], + ["rmsnorm_rht_amax_sm100"], + ) + + operation = operations[0] + torch_binding = operation.status(_REGISTRY_MODULE.FrontendTarget.TORCH) + jax_binding = operation.status(_REGISTRY_MODULE.FrontendTarget.JAX) + self.assertIsInstance(torch_binding, _REGISTRY_MODULE.TargetBinding) + self.assertIsInstance(jax_binding, _REGISTRY_MODULE.TargetBinding) + self.assertEqual( + torch_binding.qualified_name, + "cudnn.rmsnorm_rht_amax.api:rmsnorm_rht_amax_wrapper_sm100", + ) + self.assertEqual(torch_binding.target_only_parameters, ("current_stream",)) + self.assertEqual( + dict(torch_binding.output_map), + {"output": "o_tensor", "amax": "amax_tensor"}, + ) + self.assertEqual( + dict(jax_binding.output_map), + {"output": "output", "amax": "amax"}, + ) + self.assertEqual(operation.parity_case, "rmsnorm_rht_amax") + self.assertEqual(_REGISTRY_MODULE.FRONTEND_OPERATION_REGISTRY.audit(), ()) + self.assertEqual( + _REGISTRY_MODULE.FRONTEND_OPERATION_REGISTRY.audit(require_jax_complete=True), + (), + ) + + def test_optional_jax_namespace_does_not_eagerly_import_frameworks(self): + optional_roots = ("jax", "cutlass", "torch") + before = {name for name in sys.modules if name.split(".", 1)[0] in optional_roots} + + jax_namespace = importlib.import_module(f"{_TEST_PACKAGE}.jax") + + after = {name for name in sys.modules if name.split(".", 1)[0] in optional_roots} + self.assertEqual(after - before, set()) + self.assertTrue(callable(jax_namespace.rmsnorm_rht_amax_sm100)) + self.assertEqual( + jax_namespace.RmsNormRhtAmaxResult._fields, + ("output", "amax"), + ) + self.assertNotIn(f"{_TEST_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) + + def test_jax_optional_extra_is_framework_named_and_torch_free(self): + pyproject = (_REPO_ROOT / "pyproject.toml").read_text() + optional_dependencies = re.search( + r"(?ms)^\[project\.optional-dependencies\]\n(?P.*?)(?=^\[|\Z)", + pyproject, + ) + self.assertIsNotNone(optional_dependencies) + body = optional_dependencies.group("body") + self.assertRegex(body, r"(?m)^jax = \[$") + self.assertNotRegex(body, r"(?m)^cutedsl-jax = \[$") + + jax_dependencies = re.search( + r"(?ms)^jax = \[\n(?P.*?)^\]$", + body, + ) + self.assertIsNotNone(jax_dependencies) + self.assertNotIn('"torch"', jax_dependencies.group("body")) + + def test_catalog_resolves_target_bindings_lazily(self): + def torch_impl( + x_tensor, + w_tensor, + eps=1e-5, + num_threads=None, + rows_per_cta=None, + current_stream=None, + ): + return x_tensor, w_tensor, current_stream + + def jax_impl( + x, + weight, + *, + eps=1e-5, + num_threads=None, + rows_per_cta=None, + ): + return x, weight + + modules = { + "cudnn.rmsnorm_rht_amax.api": types.SimpleNamespace(rmsnorm_rht_amax_wrapper_sm100=torch_impl), + "cudnn.jax.rmsnorm_rht_amax": types.SimpleNamespace(rmsnorm_rht_amax_sm100=jax_impl), + } + operation = _FRONTEND.registered_operations()[0] + with mock.patch.object( + _REGISTRY_MODULE.importlib, + "import_module", + side_effect=lambda name: modules[name], + ): + self.assertIs( + operation.resolve(_REGISTRY_MODULE.FrontendTarget.TORCH), + torch_impl, + ) + self.assertIs( + operation.resolve(_REGISTRY_MODULE.FrontendTarget.JAX), + jax_impl, + ) + + def test_semantic_default_drift_is_rejected(self): + def drifted_jax_impl( + x, + weight, + *, + eps=1e-4, + num_threads=None, + rows_per_cta=None, + ): + return x, weight + + operation = _FRONTEND.registered_operations()[0] + with mock.patch.object( + _REGISTRY_MODULE.importlib, + "import_module", + return_value=types.SimpleNamespace(rmsnorm_rht_amax_sm100=drifted_jax_impl), + ): + with self.assertRaisesRegex(TypeError, "default.*eps.*drifted"): + operation.resolve(_REGISTRY_MODULE.FrontendTarget.JAX) + + def test_operation_ownership_cannot_overlap(self): + targets = { + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), + } + + for duplicate_kind in ("api", "kernel"): + with self.subTest(duplicate_kind=duplicate_kind): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + @_REGISTRY_MODULE.frontend_operation( + name="first", + targets=targets, + api_anchors=("example.api:First",), + kernel_anchors=("example.kernel:First.kernel",), + output_names=("output",), + parity_case="first", + registry=local_registry, + ) + def first(x): + return x + + second_api = "example.api:First" if duplicate_kind == "api" else "example.api:Second" + second_kernel = "example.kernel:First.kernel" if duplicate_kind == "kernel" else "example.kernel:Second.kernel" + with self.assertRaisesRegex(ValueError, "overlapping ownership"): + + @_REGISTRY_MODULE.frontend_operation( + name="second", + targets=targets, + api_anchors=(second_api,), + kernel_anchors=(second_kernel,), + output_names=("output",), + parity_case="second", + registry=local_registry, + ) + def second(x): + return x + + self.assertEqual(first("value"), "value") + + def test_normalized_target_names_cannot_be_declared_twice(self): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + with self.assertRaisesRegex(ValueError, "same target more than once"): + + @_REGISTRY_MODULE.frontend_operation( + name="duplicate_jax", + targets={ + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), + "JAX": _binding("example", "other_jax_impl"), + }, + api_anchors=("example.api:Example",), + kernel_anchors=("example.kernel:Example.kernel",), + output_names=("output",), + parity_case="duplicate_jax", + registry=local_registry, + ) + def duplicate_jax(x): + return x + + def test_missing_jax_target_requires_an_explicit_gap(self): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + with self.assertRaisesRegex(ValueError, "missing=.*jax"): + + @_REGISTRY_MODULE.frontend_operation( + name="missing_jax", + targets={_REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl")}, + api_anchors=("example.api:Example",), + kernel_anchors=("example.kernel:Example.kernel",), + output_names=("output",), + parity_case=None, + registry=local_registry, + ) + def missing_jax(x): + return x + + def test_torch_cannot_be_declared_as_optional_gap(self): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + with self.assertRaisesRegex(ValueError, "canonical Torch binding"): + + @_REGISTRY_MODULE.frontend_operation( + name="missing_torch", + targets={ + _REGISTRY_MODULE.FrontendTarget.TORCH: _REGISTRY_MODULE.TargetGap( + reason="not implemented", + tracking_issue="https://example.invalid/issue/1", + ), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), + }, + api_anchors=("example.api:Example",), + kernel_anchors=("example.kernel:Example.kernel",), + output_names=("output",), + parity_case="missing_torch", + registry=local_registry, + ) + def missing_torch(x): + return x + + def test_declared_jax_gap_is_visible_but_not_a_structural_failure(self): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + @_REGISTRY_MODULE.frontend_operation( + name="known_gap", + targets={ + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), + _REGISTRY_MODULE.FrontendTarget.JAX: _REGISTRY_MODULE.TargetGap( + reason="not migrated", + tracking_issue="https://example.invalid/issue/1", + ), + }, + api_anchors=("example.api:Example",), + kernel_anchors=("example.kernel:Example.kernel",), + output_names=("output",), + parity_case=None, + registry=local_registry, + ) + def known_gap(x): + return x + + self.assertEqual(known_gap("value"), "value") + self.assertEqual(local_registry.audit(), ()) + issues = local_registry.audit(require_jax_complete=True) + self.assertEqual(len(issues), 1) + self.assertIn("known_gap:jax", issues[0]) + with self.assertRaisesRegex(NotImplementedError, "not migrated"): + local_registry.get("known_gap").resolve(_REGISTRY_MODULE.FrontendTarget.JAX) + + def test_jax_binding_requires_a_parity_case(self): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + with self.assertRaisesRegex(ValueError, "no registered parity case"): + + @_REGISTRY_MODULE.frontend_operation( + name="missing_parity_case", + targets={ + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), + }, + api_anchors=("example.api:Example",), + kernel_anchors=("example.kernel:Example.kernel",), + output_names=("output",), + parity_case=None, + registry=local_registry, + ) + def missing_parity_case(x): + return x + + def test_all_discovered_python_apis_are_registered_or_baselined(self): + discovered = _discover_python_api_anchors() + registered = {anchor for operation in _FRONTEND.registered_operations() for anchor in operation.api_anchors} + gaps = set(_FRONTEND.known_jax_gaps()) + + self.assertFalse(registered & gaps) + self.assertEqual(discovered, registered | gaps) + + def test_registered_jax_gaps_cannot_grow_silently(self): + registered_gap_ids = { + operation.name + for operation in _FRONTEND.registered_operations() + if isinstance( + operation.status(_REGISTRY_MODULE.FrontendTarget.JAX), + _REGISTRY_MODULE.TargetGap, + ) + } + + self.assertEqual( + registered_gap_ids, + set(_FRONTEND.known_registered_jax_gap_ids()), + ) + + def test_all_cute_kernels_have_a_semantic_operation_owner(self): + discovered = _discover_cute_kernel_anchors() + registered = {anchor for operation in _FRONTEND.registered_operations() for anchor in operation.kernel_anchors} + gaps = set(_FRONTEND.known_kernel_ownership_gaps()) + + self.assertFalse(registered & gaps) + self.assertEqual(discovered, registered | gaps) + + def test_target_source_signatures_match_semantic_mappings(self): + operation = _FRONTEND.registered_operations()[0] + contract_path = _CUDNN_ROOT / "frontend" / "rmsnorm_rht_amax.py" + contract_parameters = _ast_function_parameters( + contract_path, + "_rmsnorm_rht_amax_sm100_contract", + ) + target_sources = { + _REGISTRY_MODULE.FrontendTarget.TORCH: ( + _CUDNN_ROOT / "rmsnorm_rht_amax" / "api.py", + "rmsnorm_rht_amax_wrapper_sm100", + ), + _REGISTRY_MODULE.FrontendTarget.JAX: ( + _CUDNN_ROOT / "jax" / "rmsnorm_rht_amax.py", + "rmsnorm_rht_amax_sm100", + ), + } + + self.assertEqual( + set(contract_parameters), + {"x", "weight", "eps", "num_threads", "rows_per_cta"}, + ) + for target, (path, function_name) in target_sources.items(): + with self.subTest(target=target.value): + binding = operation.status(target) + self.assertIsInstance(binding, _REGISTRY_MODULE.TargetBinding) + target_parameters = _ast_function_parameters(path, function_name) + expected_target_names = set(binding.parameter_map.values()) | set(binding.target_only_parameters) + self.assertEqual(set(target_parameters), expected_target_names) + for semantic_name, target_name in binding.parameter_map.items(): + self.assertEqual( + contract_parameters[semantic_name][1], + target_parameters[target_name][1], + ) + + torch_parameters = _ast_function_parameters(*target_sources[_REGISTRY_MODULE.FrontendTarget.TORCH]) + self.assertEqual( + tuple(torch_parameters), + ( + "x_tensor", + "w_tensor", + "eps", + "num_threads", + "rows_per_cta", + "current_stream", + ), + ) + self.assertTrue(all(kind == "positional" for kind, _ in torch_parameters.values())) + + jax_parameters = _ast_function_parameters(*target_sources[_REGISTRY_MODULE.FrontendTarget.JAX]) + self.assertEqual(jax_parameters["x"][0], "positional") + self.assertEqual(jax_parameters["weight"][0], "positional") + for name in ("eps", "num_threads", "rows_per_cta"): + self.assertEqual(jax_parameters[name][0], "keyword_only") + + def test_available_jax_binding_has_target_lifecycle_tests(self): + test_names_by_path = {} + for path in ( + _REPO_ROOT / "test" / "python" / "fe_api" / "test_rmsnorm_rht_amax.py", + _REPO_ROOT / "test" / "python" / "fe_api" / "test_jax_rmsnorm_rht_amax.py", + ): + tree = ast.parse(path.read_text(), filename=str(path)) + test_names_by_path[path.name] = {node.name for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))} + + for operation in _FRONTEND.registered_operations(): + if not isinstance( + operation.status(_REGISTRY_MODULE.FrontendTarget.JAX), + _REGISTRY_MODULE.TargetBinding, + ): + continue + case = operation.parity_case + self.assertIn( + f"test_{case}_wrapper", + test_names_by_path["test_rmsnorm_rht_amax.py"], + ) + self.assertIn( + f"test_jax_{case}_jit", + test_names_by_path["test_jax_rmsnorm_rht_amax.py"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_cutedsl_adapter.py b/test/python/fe_api/test_jax_cutedsl_adapter.py new file mode 100644 index 000000000..947195173 --- /dev/null +++ b/test/python/fe_api/test_jax_cutedsl_adapter.py @@ -0,0 +1,365 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CPU-only contract tests for the CuTe DSL JAX tracing adapter.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +import sys +import unittest + +try: + import pytest +except ImportError: + # Keep this contract test runnable with the standard library alone. + pass +else: + pytestmark = pytest.mark.L0 + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "jax" / "cutedsl.py" +_SPEC = importlib.util.spec_from_file_location("cudnn_jax_cutedsl_poc", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE +_SPEC.loader.exec_module(_MODULE) + +BufferInitialization = _MODULE.BufferInitialization +BufferSpec = _MODULE.BufferSpec +TensorLayout = _MODULE.TensorLayout +_call_cutedsl_with_modules = _MODULE._call_cutedsl_with_modules + + +class _Array: + def __init__(self, shape, dtype, label): + self.shape = tuple(shape) + self.dtype = dtype + self.label = label + + def __repr__(self): + return f"_Array({self.label!r}, {self.shape!r}, {self.dtype!r})" + + +class _ShapeDtypeStruct: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _FakeJax: + ShapeDtypeStruct = _ShapeDtypeStruct + + +class _FakeJaxNumpy: + def __init__(self): + self.allocations = [] + + def zeros(self, shape, dtype): + result = _Array(shape, dtype, "zeros") + self.allocations.append(("zeros", result, None)) + return result + + def full(self, shape, value, dtype): + result = _Array(shape, dtype, f"full({value})") + self.allocations.append(("full", result, value)) + return result + + +class _FakeCutlassTensorSpec: + def __init__( + self, + *, + layout, + mode, + static, + ptr_assumed_align, + divisibility, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class _FakeCutlassJax: + TensorSpec = _FakeCutlassTensorSpec + + def __init__(self): + self.calls = [] + + def cutlass_call(self, fn, **options): + self.calls.append((fn, options)) + known_options = { + "output_shape_dtype", + "input_spec", + "output_spec", + "input_output_aliases", + "allow_cuda_graph", + "compile_options", + "use_static_tensors", + } + static_args = {k: v for k, v in options.items() if k not in known_options} + + def invoke(*inputs): + aliases = options["input_output_aliases"] + alias_by_output = {output_idx: input_idx for input_idx, output_idx in aliases.items()} + results = [] + fresh_results = [] + for result_idx, metadata in enumerate(options["output_shape_dtype"]): + if result_idx in alias_by_output: + result = inputs[alias_by_output[result_idx]] + else: + result = _Array( + metadata.shape, + metadata.dtype, + f"result-{result_idx}", + ) + fresh_results.append(result) + results.append(result) + + fn("xla-stream", *inputs, *fresh_results, **static_args) + return tuple(results) + + return invoke + + +class CallCutedslAdapterTest(unittest.TestCase): + def _call(self, kernel, inputs, **kwargs): + fake_jnp = _FakeJaxNumpy() + fake_cutlass_jax = _FakeCutlassJax() + result = _call_cutedsl_with_modules( + kernel, + inputs, + jax_module=_FakeJax, + jax_numpy_module=fake_jnp, + cutlass_jax_module=fake_cutlass_jax, + **kwargs, + ) + return result, fake_jnp, fake_cutlass_jax + + def test_workspace_is_passed_to_launcher_but_hidden_from_results(self): + seen = [] + + def kernel(stream, x, output, workspace, *, scale): + seen.append((stream, x, output, workspace, scale)) + + x = _Array((8,), "f32", "x") + result, _, bridge = self._call( + kernel, + (x,), + outputs=(BufferSpec("output", (8,), "f32"),), + workspaces=(BufferSpec("workspace", (128,), "u8"),), + static_args={"scale": 2.0}, + ) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].label, "result-0") + self.assertEqual(seen[0][0], "xla-stream") + self.assertIs(seen[0][1], x) + self.assertEqual(seen[0][2].label, "result-0") + self.assertEqual(seen[0][3].label, "result-1") + self.assertEqual(seen[0][4], 2.0) + self.assertEqual(len(bridge.calls[0][1]["output_shape_dtype"]), 2) + + def test_initialized_buffers_are_inputs_aliased_to_results(self): + seen = [] + + def kernel(stream, x, output, workspace): + seen.append((output, workspace)) + + x = _Array((4,), "bf16", "x") + result, fake_jnp, bridge = self._call( + kernel, + (x,), + outputs=( + BufferSpec( + "output", + (4,), + "bf16", + initialization=BufferInitialization.VALUE, + fill_value=float("-inf"), + ), + ), + workspaces=( + BufferSpec( + "workspace", + (32,), + "u8", + initialization=BufferInitialization.ZERO, + ), + ), + ) + + aliases = bridge.calls[0][1]["input_output_aliases"] + self.assertEqual(aliases, {1: 0, 2: 1}) + self.assertEqual(bridge.calls[0][1]["input_spec"], (None, None, None)) + self.assertEqual(bridge.calls[0][1]["output_spec"], (None, None)) + self.assertIs(result[0], fake_jnp.allocations[0][1]) + self.assertIs(seen[0][0], fake_jnp.allocations[0][1]) + self.assertIs(seen[0][1], fake_jnp.allocations[1][1]) + self.assertEqual(fake_jnp.allocations[0][0], "full") + self.assertEqual(fake_jnp.allocations[1][0], "zeros") + + def test_user_alias_is_reconstructed_in_canonical_output_position(self): + seen = [] + + def kernel(stream, x, output, workspace): + seen.append((x, output, workspace)) + + x = _Array((16,), "f16", "x") + result, _, bridge = self._call( + kernel, + (x,), + outputs=(BufferSpec("output", (16,), "f16"),), + workspaces=(BufferSpec("workspace", (8,), "u8"),), + input_output_aliases={0: 0}, + ) + + self.assertIs(result[0], x) + self.assertIs(seen[0][0], x) + self.assertIs(seen[0][1], x) + self.assertEqual(seen[0][2].label, "result-1") + self.assertEqual(bridge.calls[0][1]["input_output_aliases"], {0: 0}) + + def test_middle_alias_preserves_fresh_result_order(self): + seen = [] + + def kernel(stream, x, first, aliased, third, workspace): + seen.append((first, aliased, third, workspace)) + + x = _Array((4,), "f32", "x") + result, _, _ = self._call( + kernel, + (x,), + outputs=( + BufferSpec("first", (4,), "f32"), + BufferSpec("aliased", (4,), "f32"), + BufferSpec("third", (4,), "f32"), + ), + workspaces=(BufferSpec("workspace", (16,), "u8"),), + input_output_aliases={0: 1}, + ) + + self.assertEqual([value.label for value in result], ["result-0", "x", "result-2"]) + self.assertEqual( + [value.label for value in seen[0]], + ["result-0", "x", "result-2", "result-3"], + ) + + def test_layout_metadata_is_forwarded(self): + layout = TensorLayout( + layout=(1, 0), + mode=(1, 0), + static=True, + ptr_assumed_align=128, + divisibility=(16, 8), + ) + + result, _, bridge = self._call( + lambda stream, x, output: None, + (_Array((8, 16), "f32", "x"),), + outputs=(BufferSpec("output", (8, 16), "f32", layout),), + input_layouts=(layout,), + allow_cuda_graph=False, + compile_options="--example-option", + use_static_tensors=True, + ) + + self.assertEqual(len(result), 1) + options = bridge.calls[0][1] + input_spec = options["input_spec"][0] + output_spec = options["output_spec"][0] + for spec in (input_spec, output_spec): + self.assertEqual(spec.layout, (1, 0)) + self.assertEqual(spec.mode, (1, 0)) + self.assertTrue(spec.static) + self.assertEqual(spec.ptr_assumed_align, 128) + self.assertEqual(spec.divisibility, (16, 8)) + self.assertFalse(options["allow_cuda_graph"]) + self.assertEqual(options["compile_options"], "--example-option") + self.assertTrue(options["use_static_tensors"]) + + def test_default_result_layout_uses_cutlass_inference(self): + _, _, bridge = self._call( + lambda stream, x, output: None, + (_Array((8,), "f32", "x"),), + outputs=(BufferSpec("output", (8,), "f32"),), + ) + + self.assertIsNone(bridge.calls[0][1]["output_spec"][0]) + + def test_rejects_invalid_specs_and_aliases(self): + with self.assertRaisesRegex(ValueError, "at least one public output"): + self._call(lambda stream, x: None, (_Array((1,), "f32", "x"),), outputs=()) + + with self.assertRaisesRegex(ValueError, "must be unique"): + self._call( + lambda stream, x, y, z: None, + (_Array((1,), "f32", "x"),), + outputs=(BufferSpec("same", (1,), "f32"),), + workspaces=(BufferSpec("same", (1,), "u8"),), + ) + + with self.assertRaisesRegex(ValueError, "only public outputs"): + self._call( + lambda stream, x, y, z: None, + (_Array((1,), "f32", "x"),), + outputs=(BufferSpec("output", (1,), "f32"),), + workspaces=(BufferSpec("workspace", (1,), "u8"),), + input_output_aliases={0: 1}, + ) + + with self.assertRaisesRegex(ValueError, "Expected 1 input layout"): + self._call( + lambda stream, x, y: None, + (_Array((1,), "f32", "x"),), + outputs=(BufferSpec("output", (1,), "f32"),), + input_layouts=(), + ) + + with self.assertRaisesRegex(ValueError, "shape/dtype"): + self._call( + lambda stream, x, y: None, + (_Array((2,), "f32", "x"),), + outputs=(BufferSpec("output", (1,), "f32"),), + input_output_aliases={0: 0}, + ) + + with self.assertRaisesRegex(ValueError, "identical TensorLayout"): + self._call( + lambda stream, x, y: None, + (_Array((2, 2), "f32", "x"),), + outputs=(BufferSpec("output", (2, 2), "f32"),), + input_layouts=(TensorLayout(mode=(1, 0)),), + input_output_aliases={0: 0}, + ) + + with self.assertRaisesRegex(TypeError, "flat sequence"): + self._call( + lambda stream, x, y: None, + ((_Array((1,), "f32", "nested"),),), + outputs=(BufferSpec("output", (1,), "f32"),), + ) + + with self.assertRaisesRegex(ValueError, "reserved by cutlass_call"): + self._call( + lambda stream, x, y: None, + (_Array((1,), "f32", "x"),), + outputs=(BufferSpec("output", (1,), "f32"),), + static_args={"compile_options": "not a kernel argument"}, + ) + + with self.assertRaisesRegex(ValueError, "permutation"): + self._call( + lambda stream, x, y: None, + (_Array((1, 1), "f32", "x"),), + outputs=(BufferSpec("output", (1, 1), "f32"),), + input_layouts=(TensorLayout(layout=(0, 0)),), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py new file mode 100644 index 000000000..3a609b8c5 --- /dev/null +++ b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""GPU proof-of-concept test for the JAX RMSNorm + RHT + amax wrapper.""" + +import math + +import pytest + + +def _hadamard_16(jnp): + values = [[1.0 if ((row & col).bit_count() % 2 == 0) else -1.0 for col in range(16)] for row in range(16)] + return jnp.asarray(values, dtype=jnp.float32) / math.sqrt(16) + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_abstract_shape(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + + from cudnn.jax import rmsnorm_rht_amax_sm100 + + x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) + weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) + output, amax = jax.eval_shape(rmsnorm_rht_amax_sm100, x, weight) + + assert output.shape == (256, 2048) + assert output.dtype == jnp.bfloat16 + assert amax.shape == (128,) + assert amax.dtype == jnp.float32 + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_jit(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + + gpu_devices = [device for device in jax.local_devices() if device.platform == "gpu"] + if not gpu_devices: + pytest.skip("JAX CUDA device is not available") + + capable_devices = [] + reported_capabilities = [] + for device in gpu_devices: + capability = getattr(device, "compute_capability", None) + if capability is None: + continue + major, minor = (int(value) for value in str(capability).split(".", 1)) + reported_capabilities.append(f"SM{major}{minor}") + if major * 10 + minor >= 100: + capable_devices.append(device) + if not capable_devices: + reported = ", ".join(reported_capabilities) or "unknown capabilities" + pytest.skip(f"RMSNorm + RHT requires SM100+; local GPUs report {reported}") + device = capable_devices[0] + + from cudnn.jax import rmsnorm_rht_amax_sm100 + + m, n = 256, 2048 + rows_per_cta = 2 + eps = 1e-5 + x = jax.device_put( + jax.random.normal(jax.random.key(0), (m, n), dtype=jnp.bfloat16), + device, + ) + weight = jax.device_put( + jax.random.normal(jax.random.key(1), (n,), dtype=jnp.bfloat16), + device, + ) + + @jax.jit + def run(x, weight): + return rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=128, + rows_per_cta=rows_per_cta, + ) + + lowered = run.lower(x, weight) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + output, amax = compiled(x, weight) + output.block_until_ready() + second_output, second_amax = compiled(x, weight) + second_output.block_until_ready() + + x_f32 = x.astype(jnp.float32) + normalized = x_f32 * jax.lax.rsqrt(jnp.mean(jnp.square(x_f32), axis=-1, keepdims=True) + eps) + normalized *= weight.astype(jnp.float32)[None, :] + reference = (normalized.reshape(m, n // 16, 16) @ _hadamard_16(jnp)).reshape(m, n) + amax_reference = jnp.max( + jnp.abs(reference).reshape(m // rows_per_cta, rows_per_cta, n), + axis=(1, 2), + ) + + assert output.shape == (m, n) + assert output.dtype == jnp.bfloat16 + assert amax.shape == (m // rows_per_cta,) + assert amax.dtype == jnp.float32 + assert jnp.allclose(output.astype(jnp.float32), reference, atol=4e-2, rtol=1e-2) + assert jnp.allclose(amax, amax_reference, atol=2e-3, rtol=1e-3) + assert jnp.array_equal(second_output, output) + assert jnp.array_equal(second_amax, amax) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_config.py b/test/python/fe_api/test_rmsnorm_rht_amax_config.py new file mode 100644 index 000000000..a933cdeac --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_config.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CPU-only tests for shared RMSNorm + RHT launch configuration.""" + +import importlib.util +from pathlib import Path +import sys +import unittest + +try: + import pytest +except ImportError: + # Keep this contract test runnable with the standard library alone. + pass +else: + pytestmark = pytest.mark.L0 + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_rmsnorm_rht_amax_config.py" +_SPEC = importlib.util.spec_from_file_location("cudnn_rmsnorm_rht_amax_config", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE +_SPEC.loader.exec_module(_MODULE) + + +class RmsNormRhtAmaxLaunchConfigTest(unittest.TestCase): + def test_package_does_not_eagerly_import_torch_api(self): + package_dir = _MODULE_PATH.parent / "rmsnorm_rht_amax" + module_name = "cudnn_rmsnorm_rht_amax_lazy_import_test" + package_spec = importlib.util.spec_from_file_location( + module_name, + package_dir / "__init__.py", + submodule_search_locations=[str(package_dir)], + ) + assert package_spec is not None and package_spec.loader is not None + package = importlib.util.module_from_spec(package_spec) + sys.modules[module_name] = package + try: + package_spec.loader.exec_module(package) + self.assertNotIn(f"{module_name}.api", sys.modules) + finally: + sys.modules.pop(module_name, None) + + def test_resolves_tuned_defaults(self): + self.assertEqual( + _MODULE.resolve_launch_config(256, 2048), + (128, 2), + ) + + def test_preserves_valid_overrides(self): + self.assertEqual( + _MODULE.resolve_launch_config( + 256, + 4096, + num_threads=256, + rows_per_cta=4, + ), + (256, 4), + ) + + def test_rejects_invalid_dimensions_and_launch_parameters(self): + cases = ( + ({"m": 0, "n": 2048}, "M must be positive"), + ({"m": 256, "n": 0}, "N must be positive"), + ({"m": 256, "n": 2047}, "Hadamard block size"), + ( + {"m": 256, "n": 2048, "num_threads": 512}, + "EPT=4 must be >= 8 and divisible by 8", + ), + ( + {"m": 255, "n": 2048, "rows_per_cta": 2}, + "M must be divisible", + ), + ) + for kwargs, message in cases: + with self.subTest(kwargs=kwargs): + with self.assertRaisesRegex(ValueError, message): + _MODULE.resolve_launch_config(**kwargs) + + +if __name__ == "__main__": + unittest.main() From ccf48415f001ce17ca634449cb553b04968bb90a Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 30 Jun 2026 15:51:56 -0500 Subject: [PATCH 02/33] Align Torch and JAX FE operation names --- docs/fe-oss-apis/cutedsl-jax-design.md | 58 ++++--- docs/fe-oss-apis/overview.md | 22 +-- docs/fe-oss-apis/rmsnorm_rht_amax.md | 52 ++++-- python/cudnn/README.md | 20 ++- python/cudnn/__init__.py | 1 + python/cudnn/frontend/_registry.py | 11 ++ python/cudnn/frontend/rmsnorm_rht_amax.py | 14 +- python/cudnn/rmsnorm_rht_amax/__init__.py | 1 + python/cudnn/rmsnorm_rht_amax/api.py | 28 +++- .../fe_api/test_frontend_target_parity.py | 149 ++++++++++++++---- test/python/fe_api/test_rmsnorm_rht_amax.py | 31 ++++ 11 files changed, 296 insertions(+), 91 deletions(-) diff --git a/docs/fe-oss-apis/cutedsl-jax-design.md b/docs/fe-oss-apis/cutedsl-jax-design.md index a38f4edd5..bb76cda84 100644 --- a/docs/fe-oss-apis/cutedsl-jax-design.md +++ b/docs/fe-oss-apis/cutedsl-jax-design.md @@ -143,23 +143,29 @@ Each operation is registered once with: FrontendOperationSpec name stable semantic/variant identifier contract_signature internal common parameters/defaults - targets[TORCH] required canonical wrapper binding - targets[JAX] optional binding or explicit TargetGap + targets[TORCH] required aligned high-level binding + targets[JAX] same-named binding or explicit TargetGap parameter_map semantic name -> target API name output_map semantic role -> target result name target_only_parameters e.g. Torch current_stream - api_anchors exact public class/wrapper symbols it owns + api_anchors exact legacy class/wrapper compatibility symbols kernel_anchors exact @cute.kernel symbols it owns parity_case required when a JAX binding exists ``` -The target APIs remain separate and explicit: +The target APIs remain separate and explicit while sharing an exact semantic +operation name: ```python -# Canonical existing Torch API, unchanged. +# Aligned Torch API. +from cudnn import rmsnorm_rht_amax_sm100 + +torch_result = rmsnorm_rht_amax_sm100(...) + +# Canonical existing Torch compatibility API, unchanged. from cudnn import rmsnorm_rht_amax_wrapper_sm100 -torch_result = rmsnorm_rht_amax_wrapper_sm100(...) +legacy_torch_result = rmsnorm_rht_amax_wrapper_sm100(...) # Optional JAX-native API. from cudnn.jax import rmsnorm_rht_amax_sm100 @@ -173,12 +179,13 @@ ambiguous and brittle. A `target=` argument inside the call would also become part of the traced call surface. Importing `cudnn.jax` is the explicit opt-in and provides an ordinary stable function to `jax.jit`. -Semantic parity does not mean exact Python signatures or return containers. -Torch legitimately retains `current_stream`, singleton-padding compatibility, -eager allocation, `TupleDict`, and its explicit class lifecycle. JAX owns its -stream and buffers, consumes abstract values, and returns a standard pytree. The -catalog instead checks shared logical operands, common option names/defaults, -output roles, validation/inference, and registered numerical parity coverage. +The aligned high-level functions have the same symbol, logical operand names, +common option names/defaults, and output-role names. Semantic parity still does +not require identical return container types or framework-only parameters. +Torch legitimately retains `current_stream`, eager allocation, and `TupleDict`; +JAX owns its stream and buffers and returns a standard pytree. The legacy Torch +wrapper and class also retain singleton-padding compatibility, their historical +names, and their explicit lifecycle. The registry requires both target keys. A genuinely unavailable target must be represented by `TargetGap(reason, tracking_issue)` rather than by omission. The @@ -534,18 +541,26 @@ loading before measuring capture/replay. Recommended public shape: ```python -# Existing canonical Torch API and return type remain unchanged. -from cudnn import rmsnorm_rht_amax_wrapper_sm100 +# Aligned Torch API uses the exact semantic/JAX operation name. +from cudnn import rmsnorm_rht_amax_sm100 -torch_result = rmsnorm_rht_amax_wrapper_sm100( - x_tensor=x_torch, - w_tensor=weight_torch, +torch_result = rmsnorm_rht_amax_sm100( + x_torch, + weight_torch, eps=1e-5, num_threads=128, rows_per_cta=2, current_stream=None, ) +# Existing canonical Torch compatibility API remains unchanged. +from cudnn import rmsnorm_rht_amax_wrapper_sm100 + +legacy_torch_result = rmsnorm_rht_amax_wrapper_sm100( + x_tensor=x_torch, + w_tensor=weight_torch, +) + # JAX is a separate optional install and namespace. import jax from cudnn.jax import rmsnorm_rht_amax_sm100 @@ -570,12 +585,14 @@ jax_result = run(x_jax, weight_jax) - Torch remains the default, first-class API. Its class lifecycle, wrapper signature, `TupleDict`, singleton-padding behavior, and stream controls are compatibility contracts. +- Every concrete target binding uses the exact semantic operation name. Adding + a JAX binding with a different symbol fails the dependency-free registry test. - JAX is explicitly selected by importing `cudnn.jax`; installing the base or Torch extras does not require JAX. - JAX operations use functional inputs/outputs and standard named tuples or registered pytrees. They do not emulate Torch output buffers or streams. -- Common option names and defaults should be retained when meaningful, but - exact target signatures and containers are not required. +- Common operand/option names, parameter kinds, defaults, and output roles must + match. Framework-only parameters and return-container types may differ. - Compile-affecting JAX keyword arguments are Python-static: close them over as above or name them in `jax.jit(static_argnames=...)`. Dynamic tensor/scalar operands remain explicit inputs. @@ -672,7 +689,8 @@ Status: implemented in this change. - Add the JAX CuTe call adapter. - Extract one framework-neutral launch config. - Port RMSNorm + RHT + amax. -- Add the Torch-canonical support catalog and exact-symbol ownership audit. +- Add the Torch-canonical support catalog, aligned target-symbol invariant, and + exact-symbol ownership audit. - Add CPU contract tests and an SM100 numerical test. Gate: run the GPU test on SM100, inspect lowered StableHLO for one custom call, diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index aadb46e5f..d74291e64 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -49,20 +49,22 @@ support. ## API Usage -Torch remains the canonical FE-OSS API and exposes the following wrapper and -class styles unchanged. Operations with optional JAX support are exposed in a -separate namespace: +Torch remains the canonical FE-OSS API and preserves its existing wrapper and +class styles unchanged. A JAX-enabled operation additionally exposes the same +semantic high-level function name from both framework namespaces: ```python -from cudnn import operation_name_wrapper # canonical Torch API -from cudnn.jax import operation_name # optional JAX API +from cudnn import operation_name # aligned Torch API +from cudnn.jax import operation_name # aligned optional JAX API + +from cudnn import operation_name_wrapper # existing Torch compatibility API ``` -JAX is additive: it does not replace or narrow the Torch wrapper, `TupleDict`, -class lifecycle, singleton-padding compatibility, or stream controls. The two -APIs share semantic options and kernel behavior where meaningful, while JAX -uses functional results and an XLA-owned stream. Backend selection is never -inferred from array types. +The aligned functions use the registry's exact semantic operation name and +share logical operand, option, and result-role names. Torch-only controls such +as `current_stream` remain additive. JAX does not replace or narrow the legacy +Torch wrapper, `TupleDict`, class lifecycle, singleton-padding compatibility, +or stream controls. Backend selection is never inferred from array types. ### 1. High-level wrapper diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index 030cc9ed3..3632b3285 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -8,8 +8,9 @@ This frontend integration exposes the kernel as a standard FE-OSS Python API with: - a class API (`RmsNormRhtAmaxSm100`) -- a wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) -- an optional experimental JAX API (`cudnn.jax.rmsnorm_rht_amax_sm100`) +- an aligned Torch API (`cudnn.rmsnorm_rht_amax_sm100`) +- the existing wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) +- an aligned experimental JAX API (`cudnn.jax.rmsnorm_rht_amax_sm100`) - grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse ## Shapes @@ -54,7 +55,30 @@ over every element produced by that CTA. ## API Usage -### High-level wrapper +### Aligned high-level Torch API + +The Torch and JAX namespaces expose the same semantic operation name, operand +names, option names, and result roles. Torch additionally accepts its +framework-specific stream control. + +```python +from cudnn import rmsnorm_rht_amax_sm100 + +result = rmsnorm_rht_amax_sm100( + x, + weight, + eps=1e-5, + num_threads=None, + rows_per_cta=None, + current_stream=None, +) + +output, amax = result +``` + +The aligned Torch function returns `TupleDict(output=..., amax=...)`. + +### Existing high-level wrapper ```python from cudnn import rmsnorm_rht_amax_wrapper_sm100 @@ -100,8 +124,9 @@ op.execute( ### Optional experimental JAX API -The existing Torch APIs above remain canonical and unchanged. Install the -JAX-specific optional dependencies to use the functional JAX namespace: +The existing Torch wrapper and class APIs above remain canonical and unchanged. +Install the JAX-specific optional dependencies to use the aligned functional +JAX namespace: ```bash pip install nvidia-cudnn-frontend[jax] @@ -134,8 +159,9 @@ The proof of concept requires concrete `M` and `N` during tracing and does not yet define autodiff, `vmap`, or automatic partitioning rules. `eps`, `num_threads`, and `rows_per_cta` are static compilation state; close them over as above or list them in `jax.jit(static_argnames=...)`. JAX returns a standard -`RmsNormRhtAmaxResult(output, amax)` named tuple; Torch retains its existing -`TupleDict(o_tensor, amax_tensor)` contract. See +`RmsNormRhtAmaxResult(output, amax)` named tuple. The aligned Torch function +uses the same result-role names in a `TupleDict`; the legacy wrapper retains +its existing `TupleDict(o_tensor, amax_tensor)` contract. See [CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md) for the design, workspace model, rollout plan, and current limitations. @@ -143,11 +169,11 @@ workspace model, rollout plan, and current limitations. ### Input and output tensors -- `x_tensor` / `sample_x` +- `x` / `x_tensor` / `sample_x` - Shape: `(M, N)` - Layout: row-major contiguous - Dtype: `torch.bfloat16` -- `w_tensor` / `sample_w` +- `weight` / `w_tensor` / `sample_w` - Shape: `(N,)` - Layout: contiguous - Dtype: `torch.bfloat16` @@ -171,11 +197,15 @@ workspace model, rollout plan, and current limitations. ### Wrapper return values -Returns a `TupleDict` with keys: +The aligned Torch function returns a `TupleDict` with keys: +- `output` +- `amax` + +The legacy wrapper returns a `TupleDict` with keys: - `o_tensor` - `amax_tensor` -Tuple unpacking order is `(o_tensor, amax_tensor)`. +Tuple unpacking follows each key order. ## Support surface and constraints diff --git a/python/cudnn/README.md b/python/cudnn/README.md index 7854e7af6..eb7412066 100644 --- a/python/cudnn/README.md +++ b/python/cudnn/README.md @@ -32,14 +32,17 @@ test/python/ # Test files The intended unit is a semantic operation variant, which may own one or more main/helper CuTe kernels. The existing Torch conventions remain canonical; JAX -is an optional additive namespace: +is an optional additive namespace. JAX-enabled operations expose an aligned +high-level name in both namespaces while retaining the legacy Torch wrapper: ```python -from cudnn import my_operation_wrapper -torch_result = my_operation_wrapper(torch_inputs, ...) +from cudnn import my_operation +torch_result = my_operation(torch_inputs, ...) from cudnn.jax import my_operation jax_result = my_operation(jax_inputs, ...) + +from cudnn import my_operation_wrapper # unchanged compatibility API ``` To add a new frontend-only API: @@ -50,13 +53,16 @@ To add a new frontend-only API: 2. Add the CuTe implementation. Every `@cute.kernel`, including helper kernels, must be owned by the semantic operation's exact `kernel_anchors` entry. 3. Preserve or add the canonical Torch class/wrapper API following existing - conventions. Do not replace its `TupleDict`, stream controls, output-buffer - lifecycle, or compatibility behavior to make it resemble JAX. + conventions. Add an aligned Torch high-level function whose symbol exactly + matches the semantic operation and JAX symbol. Do not replace the legacy + `TupleDict`, stream controls, output-buffer lifecycle, or compatibility + behavior to make it resemble JAX. 4. Add the functional JAX adapter using `cutlass.jax.cutlass_call`. It must infer outputs/workspace from abstract metadata, accept XLA's stream, and avoid Torch imports or host reads during tracing. -5. Register the internal semantic contract in `python/cudnn/frontend`. The Torch - binding is required; JAX must be a `TargetBinding` or explicit +5. Register the internal semantic contract in `python/cudnn/frontend`. Every + concrete target binding symbol must exactly equal the semantic operation + name. The Torch binding is required; JAX must be a `TargetBinding` or explicit `TargetGap(reason, tracking_issue)`. Record parameter/output mappings, target-only arguments, exact `api_anchors`, and exact `kernel_anchors`. 6. Add common support-domain/numerical parity cases plus Torch and JAX lifecycle diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index ccd1b4169..9a2763d35 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -274,6 +274,7 @@ def _dlopen_cudnn(): "GemmAmaxSm100": (".gemm_amax", "GemmAmaxSm100"), "gemm_amax_wrapper_sm100": (".gemm_amax", "gemm_amax_wrapper_sm100"), "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), + "rmsnorm_rht_amax_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_sm100"), "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), "grouped_gemm": (".grouped_gemm", None), "GroupedGemmSwigluSm100": (".grouped_gemm", "GroupedGemmSwigluSm100"), diff --git a/python/cudnn/frontend/_registry.py b/python/cudnn/frontend/_registry.py index 2d2cc8897..cdefb171a 100644 --- a/python/cudnn/frontend/_registry.py +++ b/python/cudnn/frontend/_registry.py @@ -158,6 +158,10 @@ def __post_init__(self) -> None: if not all(isinstance(status, (TargetBinding, TargetGap)) for status in normalized_targets.values()): raise TypeError(f"{self.name} target entries must be TargetBinding or TargetGap") + for target, status in normalized_targets.items(): + if isinstance(status, TargetBinding) and status.symbol != self.name: + raise ValueError(f"{self.name} {target.value} binding symbol must match the " f"semantic operation name; got {status.symbol!r}") + contract_names = set(self.contract_signature.parameters) output_names = set(self.output_names) for target, status in normalized_targets.items(): @@ -215,6 +219,13 @@ def _validate_binding_signature( for semantic_name, target_name in binding.parameter_map.items(): semantic_parameter = self.contract_signature.parameters[semantic_name] target_parameter = actual_parameters[target_name] + if semantic_parameter.kind != target_parameter.kind: + raise TypeError( + f"{self.name} {target.value} parameter kind for semantic " + f"parameter {semantic_name!r} drifted; expected " + f"{semantic_parameter.kind.description}, got " + f"{target_parameter.kind.description}" + ) if semantic_parameter.default != target_parameter.default: raise TypeError( f"{self.name} {target.value} default for semantic parameter " diff --git a/python/cudnn/frontend/rmsnorm_rht_amax.py b/python/cudnn/frontend/rmsnorm_rht_amax.py index 18e12a17f..49044ba73 100644 --- a/python/cudnn/frontend/rmsnorm_rht_amax.py +++ b/python/cudnn/frontend/rmsnorm_rht_amax.py @@ -27,17 +27,9 @@ targets={ FrontendTarget.TORCH: TargetBinding( module="cudnn.rmsnorm_rht_amax.api", - symbol="rmsnorm_rht_amax_wrapper_sm100", - parameter_map={ - **_COMMON_PARAMETER_NAMES, - "x": "x_tensor", - "weight": "w_tensor", - }, - output_map={ - **_COMMON_OUTPUT_NAMES, - "output": "o_tensor", - "amax": "amax_tensor", - }, + symbol="rmsnorm_rht_amax_sm100", + parameter_map=_COMMON_PARAMETER_NAMES, + output_map=_COMMON_OUTPUT_NAMES, target_only_parameters=("current_stream",), ), FrontendTarget.JAX: TargetBinding( diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 09ff4877d..3be2724da 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -14,6 +14,7 @@ "RmsNormRhtAmaxSm100", "best_num_threads", "pick_rows_per_cta", + "rmsnorm_rht_amax_sm100", "rmsnorm_rht_amax_wrapper_sm100", ] diff --git a/python/cudnn/rmsnorm_rht_amax/api.py b/python/cudnn/rmsnorm_rht_amax/api.py index 0294a3a22..9678e82a8 100644 --- a/python/cudnn/rmsnorm_rht_amax/api.py +++ b/python/cudnn/rmsnorm_rht_amax/api.py @@ -7,7 +7,6 @@ from typing import Optional from cuda.bindings import driver as cuda -import cutlass import cutlass.cute as cute import torch from cutlass import Float32 @@ -300,3 +299,30 @@ def rmsnorm_rht_amax_wrapper_sm100( ) return TupleDict(o_tensor=o_tensor, amax_tensor=amax_tensor) + + +def rmsnorm_rht_amax_sm100( + x: torch.Tensor, + weight: torch.Tensor, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + current_stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """Apply RMSNorm + RHT + amax through the aligned Torch API. + + This additive facade shares its operation name, semantic operands, options, + and result roles with ``cudnn.jax.rmsnorm_rht_amax_sm100``. The existing + ``rmsnorm_rht_amax_wrapper_sm100`` API remains unchanged for compatibility. + """ + + result = rmsnorm_rht_amax_wrapper_sm100( + x_tensor=x, + w_tensor=weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + current_stream=current_stream, + ) + return TupleDict(output=result["o_tensor"], amax=result["amax_tensor"]) diff --git a/test/python/fe_api/test_frontend_target_parity.py b/test/python/fe_api/test_frontend_target_parity.py index 87f59de1e..c3598d3fe 100644 --- a/test/python/fe_api/test_frontend_target_parity.py +++ b/test/python/fe_api/test_frontend_target_parity.py @@ -58,6 +58,13 @@ def _binding( ) +def _aligned_bindings(name): + return { + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example.torch", name), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example.jax", name), + } + + def _base_name(base): if isinstance(base, ast.Name): return base.id @@ -174,12 +181,14 @@ def test_registered_operation_keeps_torch_canonical_and_jax_optional(self): self.assertIsInstance(jax_binding, _REGISTRY_MODULE.TargetBinding) self.assertEqual( torch_binding.qualified_name, - "cudnn.rmsnorm_rht_amax.api:rmsnorm_rht_amax_wrapper_sm100", + "cudnn.rmsnorm_rht_amax.api:rmsnorm_rht_amax_sm100", ) + self.assertEqual(torch_binding.symbol, operation.name) + self.assertEqual(jax_binding.symbol, operation.name) self.assertEqual(torch_binding.target_only_parameters, ("current_stream",)) self.assertEqual( dict(torch_binding.output_map), - {"output": "o_tensor", "amax": "amax_tensor"}, + {"output": "output", "amax": "amax"}, ) self.assertEqual( dict(jax_binding.output_map), @@ -207,6 +216,35 @@ def test_optional_jax_namespace_does_not_eagerly_import_frameworks(self): ) self.assertNotIn(f"{_TEST_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) + def test_aligned_torch_name_is_exported_from_the_canonical_namespace(self): + top_level_tree = ast.parse( + (_CUDNN_ROOT / "__init__.py").read_text(), + filename=str(_CUDNN_ROOT / "__init__.py"), + ) + lazy_imports_node = next( + node.value + for node in top_level_tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_LAZY_OPTIONAL_IMPORTS" for target in node.targets) + ) + lazy_imports = ast.literal_eval(lazy_imports_node) + + operation = _FRONTEND.registered_operations()[0] + self.assertEqual( + lazy_imports[operation.name], + (".rmsnorm_rht_amax", operation.name), + ) + + package_tree = ast.parse( + (_CUDNN_ROOT / "rmsnorm_rht_amax" / "__init__.py").read_text(), + filename=str(_CUDNN_ROOT / "rmsnorm_rht_amax" / "__init__.py"), + ) + package_exports_node = next( + node.value + for node in package_tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "__all__" for target in node.targets) + ) + self.assertIn(operation.name, ast.literal_eval(package_exports_node)) + def test_jax_optional_extra_is_framework_named_and_torch_free(self): pyproject = (_REPO_ROOT / "pyproject.toml").read_text() optional_dependencies = re.search( @@ -227,14 +265,15 @@ def test_jax_optional_extra_is_framework_named_and_torch_free(self): def test_catalog_resolves_target_bindings_lazily(self): def torch_impl( - x_tensor, - w_tensor, + x, + weight, + *, eps=1e-5, num_threads=None, rows_per_cta=None, current_stream=None, ): - return x_tensor, w_tensor, current_stream + return x, weight, current_stream def jax_impl( x, @@ -247,7 +286,7 @@ def jax_impl( return x, weight modules = { - "cudnn.rmsnorm_rht_amax.api": types.SimpleNamespace(rmsnorm_rht_amax_wrapper_sm100=torch_impl), + "cudnn.rmsnorm_rht_amax.api": types.SimpleNamespace(rmsnorm_rht_amax_sm100=torch_impl), "cudnn.jax.rmsnorm_rht_amax": types.SimpleNamespace(rmsnorm_rht_amax_sm100=jax_impl), } operation = _FRONTEND.registered_operations()[0] @@ -285,19 +324,33 @@ def drifted_jax_impl( with self.assertRaisesRegex(TypeError, "default.*eps.*drifted"): operation.resolve(_REGISTRY_MODULE.FrontendTarget.JAX) - def test_operation_ownership_cannot_overlap(self): - targets = { - _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), - _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), - } + def test_semantic_parameter_kind_drift_is_rejected(self): + def drifted_jax_impl( + x, + weight, + eps=1e-5, + num_threads=None, + rows_per_cta=None, + ): + return x, weight + + operation = _FRONTEND.registered_operations()[0] + with mock.patch.object( + _REGISTRY_MODULE.importlib, + "import_module", + return_value=types.SimpleNamespace(rmsnorm_rht_amax_sm100=drifted_jax_impl), + ): + with self.assertRaisesRegex(TypeError, "parameter kind.*eps.*drifted"): + operation.resolve(_REGISTRY_MODULE.FrontendTarget.JAX) + def test_operation_ownership_cannot_overlap(self): for duplicate_kind in ("api", "kernel"): with self.subTest(duplicate_kind=duplicate_kind): local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() @_REGISTRY_MODULE.frontend_operation( name="first", - targets=targets, + targets=_aligned_bindings("first"), api_anchors=("example.api:First",), kernel_anchors=("example.kernel:First.kernel",), output_names=("output",), @@ -313,7 +366,7 @@ def first(x): @_REGISTRY_MODULE.frontend_operation( name="second", - targets=targets, + targets=_aligned_bindings("second"), api_anchors=(second_api,), kernel_anchors=(second_kernel,), output_names=("output",), @@ -333,9 +386,9 @@ def test_normalized_target_names_cannot_be_declared_twice(self): @_REGISTRY_MODULE.frontend_operation( name="duplicate_jax", targets={ - _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), - _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), - "JAX": _binding("example", "other_jax_impl"), + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example.torch", "duplicate_jax"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example.jax", "duplicate_jax"), + "JAX": _binding("example.other_jax", "duplicate_jax"), }, api_anchors=("example.api:Example",), kernel_anchors=("example.kernel:Example.kernel",), @@ -353,7 +406,7 @@ def test_missing_jax_target_requires_an_explicit_gap(self): @_REGISTRY_MODULE.frontend_operation( name="missing_jax", - targets={_REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl")}, + targets={_REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "missing_jax")}, api_anchors=("example.api:Example",), kernel_anchors=("example.kernel:Example.kernel",), output_names=("output",), @@ -375,7 +428,7 @@ def test_torch_cannot_be_declared_as_optional_gap(self): reason="not implemented", tracking_issue="https://example.invalid/issue/1", ), - _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "missing_torch"), }, api_anchors=("example.api:Example",), kernel_anchors=("example.kernel:Example.kernel",), @@ -392,7 +445,7 @@ def test_declared_jax_gap_is_visible_but_not_a_structural_failure(self): @_REGISTRY_MODULE.frontend_operation( name="known_gap", targets={ - _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "known_gap"), _REGISTRY_MODULE.FrontendTarget.JAX: _REGISTRY_MODULE.TargetGap( reason="not migrated", tracking_issue="https://example.invalid/issue/1", @@ -423,8 +476,8 @@ def test_jax_binding_requires_a_parity_case(self): @_REGISTRY_MODULE.frontend_operation( name="missing_parity_case", targets={ - _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example", "torch_impl"), - _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example", "jax_impl"), + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example.torch", "missing_parity_case"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example.jax", "missing_parity_case"), }, api_anchors=("example.api:Example",), kernel_anchors=("example.kernel:Example.kernel",), @@ -435,6 +488,26 @@ def test_jax_binding_requires_a_parity_case(self): def missing_parity_case(x): return x + def test_target_symbols_must_match_the_semantic_operation_name(self): + local_registry = _REGISTRY_MODULE.FrontendOperationRegistry() + + with self.assertRaisesRegex(ValueError, "binding symbol must match"): + + @_REGISTRY_MODULE.frontend_operation( + name="aligned_name", + targets={ + _REGISTRY_MODULE.FrontendTarget.TORCH: _binding("example.torch", "legacy_wrapper_name"), + _REGISTRY_MODULE.FrontendTarget.JAX: _binding("example.jax", "aligned_name"), + }, + api_anchors=("example.api:Example",), + kernel_anchors=("example.kernel:Example.kernel",), + output_names=("output",), + parity_case="aligned_name", + registry=local_registry, + ) + def aligned_name(x): + return x + def test_all_discovered_python_apis_are_registered_or_baselined(self): discovered = _discover_python_api_anchors() registered = {anchor for operation in _FRONTEND.registered_operations() for anchor in operation.api_anchors} @@ -476,7 +549,7 @@ def test_target_source_signatures_match_semantic_mappings(self): target_sources = { _REGISTRY_MODULE.FrontendTarget.TORCH: ( _CUDNN_ROOT / "rmsnorm_rht_amax" / "api.py", - "rmsnorm_rht_amax_wrapper_sm100", + "rmsnorm_rht_amax_sm100", ), _REGISTRY_MODULE.FrontendTarget.JAX: ( _CUDNN_ROOT / "jax" / "rmsnorm_rht_amax.py", @@ -505,21 +578,31 @@ def test_target_source_signatures_match_semantic_mappings(self): self.assertEqual( tuple(torch_parameters), ( - "x_tensor", - "w_tensor", + "x", + "weight", "eps", "num_threads", "rows_per_cta", "current_stream", ), ) - self.assertTrue(all(kind == "positional" for kind, _ in torch_parameters.values())) - - jax_parameters = _ast_function_parameters(*target_sources[_REGISTRY_MODULE.FrontendTarget.JAX]) - self.assertEqual(jax_parameters["x"][0], "positional") - self.assertEqual(jax_parameters["weight"][0], "positional") - for name in ("eps", "num_threads", "rows_per_cta"): - self.assertEqual(jax_parameters[name][0], "keyword_only") + for target in _REGISTRY_MODULE.FrontendTarget: + target_parameters = _ast_function_parameters(*target_sources[target]) + self.assertEqual(target_parameters["x"][0], "positional") + self.assertEqual(target_parameters["weight"][0], "positional") + for name in ("eps", "num_threads", "rows_per_cta"): + self.assertEqual(target_parameters[name][0], "keyword_only") + self.assertEqual(torch_parameters["current_stream"][0], "keyword_only") + + legacy_parameters = _ast_function_parameters( + _CUDNN_ROOT / "rmsnorm_rht_amax" / "api.py", + "rmsnorm_rht_amax_wrapper_sm100", + ) + self.assertEqual( + tuple(legacy_parameters), + ("x_tensor", "w_tensor", "eps", "num_threads", "rows_per_cta", "current_stream"), + ) + self.assertTrue(all(kind == "positional" for kind, _ in legacy_parameters.values())) def test_available_jax_binding_has_target_lifecycle_tests(self): test_names_by_path = {} @@ -541,6 +624,10 @@ def test_available_jax_binding_has_target_lifecycle_tests(self): f"test_{case}_wrapper", test_names_by_path["test_rmsnorm_rht_amax.py"], ) + self.assertIn( + f"test_{case}_aligned_api", + test_names_by_path["test_rmsnorm_rht_amax.py"], + ) self.assertIn( f"test_jax_{case}_jit", test_names_by_path["test_jax_rmsnorm_rht_amax.py"], diff --git a/test/python/fe_api/test_rmsnorm_rht_amax.py b/test/python/fe_api/test_rmsnorm_rht_amax.py index c00d04b02..2e6eb093a 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax.py @@ -122,3 +122,34 @@ def test_rmsnorm_rht_amax_wrapper(n, num_threads, rows_per_cta, request): assert outputs["o_tensor"].shape == (m, n) assert outputs["amax_tensor"].shape == (m // rows_per_cta,) _assert_ref_close(x, w, outputs["o_tensor"], outputs["amax_tensor"], eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_rmsnorm_rht_amax_aligned_api(request): + try: + from cudnn import rmsnorm_rht_amax_sm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + skip_ref = request.config.getoption("--skip-ref", default=False) + eps = 1e-5 + m, n = 256, 2048 + rows_per_cta = 2 + x, weight = _make_inputs(m=m, n=n) + + try: + result = rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=128, + rows_per_cta=rows_per_cta, + ) + except (ValueError, RuntimeError) as exc: + pytest.skip(f"Unsupported testcase: {exc}") + + assert tuple(result.keys()) == ("output", "amax") + assert result["output"].shape == (m, n) + assert result["amax"].shape == (m // rows_per_cta,) + _assert_ref_close(x, weight, result["output"], result["amax"], eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) From 713841fbb3a449cc8a81b2fb3b68d4d9deb7595c Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 30 Jun 2026 16:00:31 -0500 Subject: [PATCH 03/33] Add Torch and JAX documentation tabs --- docs/fe-oss-apis/rmsnorm_rht_amax.md | 134 +++++++++++------- .../fe_api/test_frontend_target_parity.py | 41 ++++++ 2 files changed, 122 insertions(+), 53 deletions(-) diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index 3632b3285..ea421fac1 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -55,12 +55,25 @@ over every element produced by that CTA. ## API Usage -### Aligned high-level Torch API +### Aligned high-level API The Torch and JAX namespaces expose the same semantic operation name, operand names, option names, and result roles. Torch additionally accepts its framework-specific stream control. +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + +Install the CuTe DSL dependencies: + +```bash +pip install nvidia-cudnn-frontend[cutedsl] +``` + ```python from cudnn import rmsnorm_rht_amax_sm100 @@ -78,7 +91,62 @@ output, amax = result The aligned Torch function returns `TupleDict(output=..., amax=...)`. -### Existing high-level wrapper +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX-specific optional dependencies: + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +The JAX optional dependency set requires Python 3.11 or newer. + +```python +import jax +from cudnn.jax import rmsnorm_rht_amax_sm100 + +eps = 1e-5 +num_threads = 128 +rows_per_cta = 2 + +@jax.jit +def run(x, weight): + return rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + +output, amax = run(x, weight) +``` + +The JAX function returns +`RmsNormRhtAmaxResult(output=..., amax=...)`, a standard named tuple and JAX +pytree. + +````` + +`````` + +The JAX proof of concept requires concrete `M` and `N` during tracing and does +not yet define autodiff, `vmap`, or automatic partitioning rules. `eps`, +`num_threads`, and `rows_per_cta` are static compilation state; close them over +as above or list them in `jax.jit(static_argnames=...)`. JAX always uses XLA's +runtime stream; Torch optionally accepts `current_stream`. See +[CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md) for the design, +workspace model, rollout plan, and current limitations. + +### Additional PyTorch APIs + +The aligned function above is additive. The existing PyTorch wrapper and class +APIs remain first class and retain their original contracts. + +#### Existing high-level wrapper ```python from cudnn import rmsnorm_rht_amax_wrapper_sm100 @@ -97,7 +165,7 @@ o_tensor, amax_tensor = result When no overrides are supplied, the wrapper uses the upstream-tuned thread table when available and an upstream-style `rows_per_cta` heuristic. -### Class API +#### Class API ```python from cudnn import RmsNormRhtAmaxSm100 @@ -122,49 +190,6 @@ op.execute( ) ``` -### Optional experimental JAX API - -The existing Torch wrapper and class APIs above remain canonical and unchanged. -Install the JAX-specific optional dependencies to use the aligned functional -JAX namespace: - -```bash -pip install nvidia-cudnn-frontend[jax] -``` - -The JAX optional dependency set requires Python 3.11 or newer. - -```python -import jax -from cudnn.jax import rmsnorm_rht_amax_sm100 - -eps = 1e-5 -num_threads = 128 -rows_per_cta = 2 - -@jax.jit -def run(x, weight): - return rmsnorm_rht_amax_sm100( - x, - weight, - eps=eps, - num_threads=num_threads, - rows_per_cta=rows_per_cta, - ) - -o, amax = run(x, weight) -``` - -The proof of concept requires concrete `M` and `N` during tracing and does not -yet define autodiff, `vmap`, or automatic partitioning rules. `eps`, -`num_threads`, and `rows_per_cta` are static compilation state; close them over -as above or list them in `jax.jit(static_argnames=...)`. JAX returns a standard -`RmsNormRhtAmaxResult(output, amax)` named tuple. The aligned Torch function -uses the same result-role names in a `TupleDict`; the legacy wrapper retains -its existing `TupleDict(o_tensor, amax_tensor)` contract. See -[CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md) for the design, -workspace model, rollout plan, and current limitations. - ## Parameters ### Input and output tensors @@ -172,18 +197,18 @@ workspace model, rollout plan, and current limitations. - `x` / `x_tensor` / `sample_x` - Shape: `(M, N)` - Layout: row-major contiguous - - Dtype: `torch.bfloat16` + - Dtype: bfloat16 (`torch.bfloat16` or `jax.numpy.bfloat16`) - `weight` / `w_tensor` / `sample_w` - Shape: `(N,)` - Layout: contiguous - - Dtype: `torch.bfloat16` -- `o_tensor` / `sample_o` + - Dtype: bfloat16 (`torch.bfloat16` or `jax.numpy.bfloat16`) +- `output` / `o_tensor` / `sample_o` - Shape: `(M, N)` - Layout: row-major contiguous - - Dtype: `torch.bfloat16` -- `amax_tensor` / `sample_amax` + - Dtype: bfloat16 (`torch.bfloat16` or `jax.numpy.bfloat16`) +- `amax` / `amax_tensor` / `sample_amax` - Shape: `(M / rows_per_cta,)` - - Dtype: `torch.float32` + - Dtype: float32 (`torch.float32` or `jax.numpy.float32`) ### Common parameters @@ -195,12 +220,15 @@ workspace model, rollout plan, and current limitations. - Rows processed by each CTA. If omitted, the wrapper uses the upstream-style heuristic over `{2, 4, 8}`. - Torch-only CUDA stream (`current_stream`); JAX always uses XLA's stream -### Wrapper return values +### Return values The aligned Torch function returns a `TupleDict` with keys: - `output` - `amax` +The aligned JAX function returns an `RmsNormRhtAmaxResult` named tuple with +fields `output` and `amax`. + The legacy wrapper returns a `TupleDict` with keys: - `o_tensor` - `amax_tensor` diff --git a/test/python/fe_api/test_frontend_target_parity.py b/test/python/fe_api/test_frontend_target_parity.py index c3598d3fe..a38e0bc90 100644 --- a/test/python/fe_api/test_frontend_target_parity.py +++ b/test/python/fe_api/test_frontend_target_parity.py @@ -25,8 +25,11 @@ _REPO_ROOT = Path(__file__).resolve().parents[3] _PYTHON_ROOT = _REPO_ROOT / "python" _CUDNN_ROOT = _PYTHON_ROOT / "cudnn" +_FE_OSS_DOCS_ROOT = _REPO_ROOT / "docs" / "fe-oss-apis" _TEST_PACKAGE = "cudnn_frontend_parity_contract_test" _REQUIRED_DEFAULT = "" +_TAB_SET_PATTERN = re.compile(r"(?ms)^(?P`{4,})\{tab-set\}\n(?P.*?)^(?P=fence)$") +_TAB_ITEM_PATTERN = re.compile(r"(?ms)^(?P`{3,})\{tab-item\} (?P