From 17bbbbfab129f21af2e164b3b563213d9533f724 Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Mon, 17 Aug 2026 13:13:00 -0700 Subject: [PATCH 1/7] Add documentation and agent skill for integrating cuDNN Frontend kernels in pure JAX workflows --- AGENTS.md | 15 ++ docs/agent-skills/jax-cudnn-frontend/SKILL.md | 175 +++++++++++++++++ .../references/contract-discovery.md | 99 ++++++++++ .../references/debugging.md | 146 ++++++++++++++ .../references/environment-discovery.md | 71 +++++++ .../references/jax-integration.md | 137 +++++++++++++ .../references/validation.md | 106 ++++++++++ .../references/worked-example.md | 94 +++++++++ .../scripts/contract_report.py | 184 ++++++++++++++++++ .../scripts/inspect_environment.py | 155 +++++++++++++++ 10 files changed, 1182 insertions(+) create mode 100644 AGENTS.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/SKILL.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/references/contract-discovery.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/references/debugging.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/references/environment-discovery.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/references/validation.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/references/worked-example.md create mode 100644 docs/agent-skills/jax-cudnn-frontend/scripts/contract_report.py create mode 100644 docs/agent-skills/jax-cudnn-frontend/scripts/inspect_environment.py diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..c8ebe4385 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,15 @@ +# Agent skills + +[`docs/agent-skills/`](docs/agent-skills/) holds procedural skills for AI +coding agents working with JAX. They aren't about navigating this repo's own +code — they cover general JAX integration work (e.g. wiring up an external +kernel library) and are stored here purely for discoverability. Each skill is +self-contained (instructions + reference docs + scripts) and tool-agnostic — +read the skill's own file directly regardless of which agent or IDE you're +using. + +- [`jax-cudnn-frontend`](docs/agent-skills/jax-cudnn-frontend/SKILL.md) — + implementing, integrating, debugging, and validating cuDNN Frontend / CuTe + DSL kernels from pure JAX. Use before wrapping, calling, porting, or + autotuning a cudnn-frontend kernel (attention variants, GEMM fusions, + experimental csrc kernel classes) from JAX. diff --git a/docs/agent-skills/jax-cudnn-frontend/SKILL.md b/docs/agent-skills/jax-cudnn-frontend/SKILL.md new file mode 100644 index 000000000..6e5d49d8b --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/SKILL.md @@ -0,0 +1,175 @@ +--- +name: jax-cudnn-frontend +description: Implement, integrate, debug, and validate cuDNN Frontend / CuTe DSL kernels from pure JAX (no PyTorch). Use when asked to wrap, call, port, or autotune a cudnn-frontend kernel — attention variants, GEMM fusions, or experimental csrc kernel classes — from JAX. Discovery-first — derives APIs and tensor contracts from the installed environment instead of trusting memory or online examples. +--- + +# cuDNN Frontend kernels from pure JAX + +You are integrating experimental GPU kernels whose APIs, tensor contracts, and +supported configurations change between package versions. **Treat everything +you think you know about these APIs as stale.** Every claim about a signature, +shape, dtype, layout, or flag must be re-derived from the installed +environment before you build on it. This skill is the procedure for doing +that, plus the failure modes that cost days when the procedure was skipped. + +## Non-negotiable rules + +1. **Never call a kernel whose contract you have not read from installed + source.** Kernel names, docs pages, GitHub examples, and your prior + knowledge describe *some* version — not necessarily the installed one. +2. **Buffer shapes come from the vendor's own allocations**, not from + signature comments or docs. Find where the installed package allocates + each buffer before calling the kernel itself (see rule of the Rosetta + stone, Phase 2). An under-allocated output is a silent buffer overflow + that corrupts *neighboring* tensors and surfaces as unrelated crashes + later — memcheck cannot see it. +3. **Compilation and execution are not correctness.** A kernel can launch, + return, and pass `block_until_ready()` while producing garbage, writing + nothing, or corrupting memory that only a later op trips over. Only value + checks against an independent reference count as passing. +4. **No PyTorch** in the implementation path unless the user asks for it. + PyTorch is permitted only as an optional cross-check oracle, clearly + labeled. +5. **Fresh process per meaningful measurement.** After any illegal memory + access or suspect result, the CUDA context and all session state are + untrustworthy — including results that *look* fine. +6. **One variable per experiment.** Build an env-var-toggled repro script + early (Phase 5) and bisect with it; never change two things between runs. + +## Workflow + +### Phase 1 — Environment discovery +Run `scripts/inspect_environment.py`. Record: JAX/jaxlib/plugin versions and +the source-commit suffix, cudnn-frontend and cutlass-dsl versions, CUDA +toolkit, GPU name + compute capability + driver, and whether the target +module/kernel is importable. Details and container-identity pitfalls (rolling +tags, arm64 lag): `references/environment-discovery.md`. + +**Exit criteria:** a pasted environment block in your working notes, and a +confirmed compute capability — kernel availability and *which kernel class +you must use* are usually arch-gated (e.g. separate `sm90_*` / `sm100_*` +classes with different constructors and argument orders). + +### Phase 2 — Locate the kernel and its Rosetta stone +Find (a) the kernel class/function in the installed package and (b) **the +vendor's own orchestration code that calls it** — typically an +`_interface.py`, `api.py`, wrapper, or test inside the installed package. +That call site is the single most valuable artifact you will find: it shows +the true argument order, every buffer allocation with exact shapes and +dtypes, workspace construction, layout transforms, and flag couplings the +constructor won't tell you about. Use `scripts/contract_report.py` to dump +signatures and call sites automatically. Procedure and grep recipes: +`references/contract-discovery.md`. + +### Phase 3 — Extract the data contract +Fill in the full checklist in `references/contract-discovery.md` (inputs, +outputs, shapes, dtypes, layout/stride requirements, alignment, scalars, +workspace, zero-init requirements, aliasing, optional-vs-required per arch, +arch constraints, coupled flags). Evidence hierarchy when sources disagree: + +1. **Executable code in the installed package** — assertions in the kernel + body (`assert`, `check_dim`, dtype checks) and the orchestration's buffer + allocations. These are the contract. +2. Tests/examples shipped *inside the installed package*. +3. Upstream source at the exact installed version (match the commit). +4. Version-labeled docs. +5. General docs, online examples, and your prior knowledge — hypothesis + generators only, never evidence. + +When (1) contradicts anything else — including signature *comments* in the +same file — (1) wins. Record the discrepancy; it often marks an API +transition and predicts other differences. When two pieces of executable code +disagree (wrapper vs kernel), the kernel body governs the kernel's tensor +contract and the wrapper governs orchestration (buffer shapes, call order). + +### Phase 4 — Choose and verify the integration mechanism +Discover what the installed bridge actually provides — do not assume: +introspect `cutlass.jax` exports and read the installed bridge source for the +launcher convention, output allocation, aliasing semantics, and spec types. +Keep **two invocation paths** implemented throughout the project: the +JAX-native path (e.g. `cutlass_call`) *and* the direct path (e.g. +`cute.compile`). The A/B between them is your primary tool for separating +bridge bugs from kernel bugs from your bugs. Mechanics, conventions, and +version-sensitive behaviors to verify: `references/jax-integration.md`. + +### Phase 5 — Minimal standalone proof +Before any notebook or abstraction: a single plain-Python script that runs +the kernel once with fixed seeds, prints **value-based** evidence (NaN +counts, checksums, fixed-position samples), and exposes every configuration +choice as an env-var toggle. This script is simultaneously your repro for +bug reports and your bisection harness. Allocate outputs with `jnp.empty` so +unwritten regions show up as NaN — a free tile-coverage diagnostic. Then run +it **N≥3 iterations** in one process (repeated execution is where buffer and +lifetime bugs hide) and **twice in fresh processes** (identical checksums; +drift is a red flag). + +### Phase 6 — JAX wrapper +Only after Phase 5 passes: wrap with the JAX-native mechanism, jit it, +confirm the jaxpr contains a custom call, and re-run the Phase 5 checks +through the wrapper — outputs should be bit-identical to the direct path. +Autodiff is a separate deliverable: `jax.custom_vjp` with the corresponding +backward kernel, whose contract gets its own full Phase 2–5 pass. Do not +assume the backward's buffers mirror the forward's. + +### Phase 7 — Validation gate +No result is "working" until it passes the gates in +`references/validation.md`: independent mathematical reference (host-side if +GPU-side references share failure modes with the kernel), analytic +invariants, repeated-execution stability, size scaling beyond the toy config +(some bugs only appear past hardware occupancy thresholds — test a size +where total tiles exceed one CTA wave), and — for anything exposing config +knobs or feeding an autotuner — a **per-configuration correctness check**, +because invalid configs can be silently wrong *and faster*. + +### Phase 8 — Debugging +Symptom-indexed decision trees from real failures: +`references/debugging.md`. Headline discipline: crashes surface at innocent +*later* operations (async execution), so the faulting op named in a traceback +is usually the victim, not the culprit. Establish the earliest corrupted +artifact, then work backwards. + +## Agent traps (each of these was committed or nearly committed) + +- Copying an invocation from docs/examples without version-matching it. +- Inferring a tensor's shape from its name (`lse` ≠ "per-block" just because + the op is block-sparse — it was per-token). +- Trusting a signature comment over the code three lines below it. +- Assuming a Python wrapper's simplified API equals the kernel's contract. +- Assuming the kernel exists / behaves the same on every GPU arch. +- Assuming default constructor arguments are valid combinations (a default + flag pair produced silent 40%-wrong output). +- Treating a clean `compute-sanitizer` run as proof of memory safety (it is + blind to overflows into valid neighboring allocations and to unwritten + output). +- Treating `block_until_ready()` or printed shapes as correctness. +- Reading device data with `np.asarray` twice (JAX caches the first host + copy; use `jnp.copy(x)` to force a fresh device read). +- Testing only at small sizes (waves-of-work bugs hide below occupancy). +- Blaming the platform/container/driver before exhausting caller-side bugs — + and, symmetrically, burning days on caller-side theories without A/B-ing + the invocation paths. +- Letting a rolling container tag define your environment. +- Silently importing torch because the vendor's public API takes torch + tensors — read one level deeper; the kernel classes underneath are + framework-agnostic via DLPack. + +## Design notes + +**Not hard-coded, on purpose:** concrete signatures, argument orders, buffer +shapes, spec/aliasing semantics, flag names, and version numbers. All of +these changed at least once during the work this skill distills, sometimes +between adjacent minor versions. Where the references show concrete code, it +is labeled *Example (version-specific)* and paired with the discovery step +that regenerates it for any version. + +**How this adapts to change:** every phase's output is derived from the +installed artifacts (signatures, assertions, vendor allocations, bridge +source), so a renamed argument or reshaped buffer changes the *result* of the +procedure, not the procedure. The scripts print what exists rather than +checking against expectations. + +**Still requires human judgment:** deciding whether an observed misbehavior +is a vendor bug worth reporting versus an undocumented contract you must +satisfy; choosing performance configs after correctness; prioritizing which +arch/version combinations to validate; and anything requiring contact with +kernel owners (known-issue status, backports, roadmap). diff --git a/docs/agent-skills/jax-cudnn-frontend/references/contract-discovery.md b/docs/agent-skills/jax-cudnn-frontend/references/contract-discovery.md new file mode 100644 index 000000000..6b912f954 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/references/contract-discovery.md @@ -0,0 +1,99 @@ +# Kernel data-contract discovery + +Goal: a written contract for the kernel covering every row of the checklist +below, with each entry traceable to installed source. `scripts/contract_report.py` +automates the mechanical parts (signatures, assertion lines, call sites). + +## No repository clone required + +This entire workflow operates on **installed packages in site-packages** — +the cuDNN Frontend CuTe DSL kernels and the cutlass-dsl JAX bridge ship as +Python source inside their pip wheels, fully introspectable via `importlib`. +Do not clone repositories by default. Two exceptions: + +- **Compiled components** (e.g. the classic cuDNN graph API's C++ backend + `.so`): the Python surface is still inspectable, but implementation + contracts must come from signatures plus empirical probing (Phase 5). +- **Tests/examples missing from the wheel**: fall to evidence tier 3 — + fetch upstream source *pinned to the exact installed version* (match the + version string or the `+commit` hash), never repository HEAD. + +## The Rosetta stone rule + +The highest-value artifact is **the vendor's own call site** for the kernel — +the orchestration module inside the installed package (`_interface.py`, +`api.py`, a `*_wrapper`, or the package's tests). Find it by grepping the +installed tree for the kernel class name: + +```bash +PKG=$(python -c "import cudnn, os; print(os.path.dirname(cudnn.__file__))") +grep -rn "KernelClassName" "$PKG" --include='*.py' | grep -v "class KernelClassName" +``` + +From the call site, extract — in this order of importance: + +1. **Buffer allocations** (`torch.empty(...)`, `np.zeros(...)`, etc. — the + framework used by the vendor's wrapper is irrelevant; the *shapes* are the + contract). This is where the single most expensive mistake in the source + experience lived: an output the operation's docs described as + per-block `(B, H, S/64)` was actually allocated per-token `(B, H, S)` by + the vendor. The 64× under-allocation corrupted neighboring buffers and + produced a week of misdirected debugging across three machines. +2. **The exact positional argument order** of the kernel call, including + `None`s for optional slots. +3. **Workspace construction** — helper functions often carry layout comments + the kernel file lacks (e.g. "fields are laid out field-major across all + B·H entries; zero the accumulator tail on the *flattened* view"). +4. **Flag couplings** — assignments like `flag_a = flag_b` in the wrapper + mean the vendor never exercises the decoupled combinations; decoupled + combinations may be silently broken (one such default combination + dropped 40% of output tiles). +5. **Layout transforms** applied to tensors before the call (transposes, + `mark_layout_dynamic`, alignment hints) — replicate their *semantics*. + +## Contract checklist + +For each kernel, fill in every row. "Unknown" is acceptable only with a plan +to determine it empirically in Phase 5. + +| Item | Where the truth lives | +|---|---| +| Compile-time config (ctor args) | `inspect.signature(Cls.__init__)` — **per architecture**; sibling arch classes (`sm90_*` vs `sm100_*`) frequently differ in both parameters and semantics (one takes `dtype` at construction, another infers it at runtime) | +| Runtime args + order | `__call__` signature; confirm against the vendor call site | +| Input shapes/ranks | Body assertions first, vendor allocations second, comments last | +| Output shapes | **Vendor allocations only.** Never from the op name or docs | +| Dtypes | Body checks (`element_type != ...`, dtype asserts) | +| Layout/stride requirements | Body assertions of the form "stride at position k must be 1" (`check_dim`-style). Positional comments in signatures can be wrong — one kernel's comment described dim order `(d, s, h, b)` while the code's reshape+assert demanded the contiguous dim at position 1, i.e. effectively `(s, d, h, b)`. **The assertion + the reshape lines are the contract; the comment is a rumor.** | +| Alignment | Vendor conversion helpers (`assumed_align=` arguments) | +| Scalar params | Signature types (e.g. runtime `Int32` vs Python-int constexpr); note constraints in asserts ("even, >= 2") | +| Workspace | Vendor helper: exact element count formula, dtype, which regions must be **zero-initialized** and on what view | +| Zero-init requirements | Any buffer the kernel accumulates into or counts with. JAX-side outputs are **uninitialized** memory — this must be handled explicitly (see jax-integration.md) | +| Aliasing/mutability | Which args the kernel writes; whether in-place semantics are expected | +| Optional vs required | `Optional[...]` in one arch's signature may be **required** in another (passing `None` produced `'NoneType' object is not subscriptable` *at trace time* on the arch where it was required) | +| Arch constraints | Grep asserts for `head_dim`, block sizes, dtype ("bwd only supports bfloat16", "requires head_dim=128") — these differ per arch and per fwd/bwd | +| Coupled flags | Wrapper assignments and heuristic-chooser functions | + +## Reading order for a kernel file + +1. `__init__` — compile-time specialization surface. +2. `__call__` signature — runtime surface. Read comments *skeptically*. +3. First ~50 lines of the `__call__` body — this is where tensors get + reshaped/remapped and asserted. The remap code defines what each position + means, overriding the signature comment. +4. Grep the body for `assert`, `check_dim`, `element_type`, `const_expr(` — + collect every constraint. +5. Count internal `.launch(` calls — multi-launch bodies historically + stress integration bridges differently than single-launch ones. + +## When name-based intuition tempts you + +Do not infer: +- shape from a tensor's name (`lse`, `stats`, `counts` — resolution and rank + vary by version and kernel); +- rank from the public API (a wrapper accepting rank-1 metadata may expand it + to rank-3 before the kernel; the kernel may index it as 3-D); +- semantics from a sibling arch's kernel (argument orders differed across + all three arch variants of the same operation in the source experience — + including where the softmax scale sits relative to the metadata tensors). + +Every one of these intuitions was wrong at least once. diff --git a/docs/agent-skills/jax-cudnn-frontend/references/debugging.md b/docs/agent-skills/jax-cudnn-frontend/references/debugging.md new file mode 100644 index 000000000..7ecbd8f11 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/references/debugging.md @@ -0,0 +1,146 @@ +# Debugging decision trees + +Symptom-indexed. Each tree is ordered: check the cheap/likely causes before +expensive theories. Overarching law: **CUDA execution is asynchronous — the +operation named in a crash is usually the victim, not the culprit.** Establish +the earliest corrupted artifact, then walk backwards. + +## Universal first moves (any GPU fault) + +1. Restart the process. Everything after an illegal access is untrustworthy. +2. Reproduce in the standalone script (fresh process), not the notebook. +3. Re-run with `CUDA_LAUNCH_BLOCKING=1` — serializes launches so errors are + attributed to the actual faulting launch. +4. If still ambiguous: `compute-sanitizer --tool memcheck python repro.py`. + Remember its blind spots: overflows into valid allocations and unwritten + output are invisible; a clean run does NOT clear the kernel. + +## Symptom: illegal memory access reported at an unrelated op (e.g. a reduce/convert kernel) + +The named op tripped over corruption left by an earlier kernel. +1. **Check every output/workspace buffer size against the vendor's own + allocations** (contract-discovery.md, Rosetta stone). The #1 historical + cause: an output allocated smaller than the kernel writes → overflow into + neighboring live buffers. Costed the most time of any bug class. +2. Check metadata tensors' *ranks/shapes* against what the kernel indexes + (a kernel indexing `t[a, b, c]` needs rank ≥ 3 — a rank-1 tensor "works" + under some compilers until it reads garbage). +3. Only then consider bridge/platform theories. + +Audit tool — is a live buffer being clobbered? + +```python +before = np.asarray(jnp.copy(suspect)) # fresh device read +run_kernel(...); out.block_until_ready() +after = np.asarray(jnp.copy(suspect)) +print(np.array_equal(before, after)) +# also compare device-side vs host-cached views: +print(int(jnp.max(suspect))) # device compute +``` + +If a supposedly-int tensor contains float bit patterns, reinterpret them +(`arr.view(np.float32)`) and fingerprint against live float tensors — a match +identifies exactly *which* buffer is being written over it. + +## Symptom: first execution correct, second execution faults + +1. First execution overflowed a buffer (see above) — sizes vs vendor + allocations, again. +2. Bridge aliasing (`input_output_aliases`) bug in the installed version — + A/B against the direct path; try the zero-prologue pattern instead of + aliasing. +3. Borrowed DLPack memory freed by its owner — `.copy()` anything returned + across a framework boundary. + +## Symptom: outputs silently all-zero / empty, no error + +1. Bridge aliasing machinery broken in installed version (a real, known bug + in one release: aliased configurations returned empty outputs). A/B + against the direct path immediately — if direct is correct, it's the + bridge; check for a newer bridge release, use aliased-output-last or the + zero-prologue workaround. +2. A buffer the kernel needed zero-initialized wasn't (counters read garbage + / kernel exited early). +3. Kernel gated off for this arch/dtype and exiting silently. + +## Symptom: structured NaN in part of the output; early regions fine + +Tiles/work-units never written. Use the coverage map (validation.md Gate 3). +1. Scheduler/config flags: check the vendor wrapper for **coupled flags** + (e.g. persistence paired with a companion scheduler). A constructor-default + combination the vendor never exercises may drop all work past one hardware + wave — test above and below `work_items == SM_count`. +2. Grid/tile mapping mismatch with your problem shape. + +## Symptom: values uniformly wrong (finite, plausible magnitude, everywhere) + +The kernel computed *something else* — usually attention/gather over the +wrong data. +1. Metadata view/permutation wrong (mode/layout specs) — kernel reads the + right buffer with the wrong coordinate mapping. +2. Metadata buffer corrupted before the kernel read it (audit tool above; + also re-upload fresh: `jnp.asarray(np.asarray(jnp.copy(meta)))`). +3. Your reference is wrong — cross-validate two independent references on + CPU before convicting the kernel (this exact dispute happened; the + reference won). + +## Symptom: "dim must be contiguous in mode k" (or similar stride assert) + +The kernel's layout contract. Read the assert AND the tensor-remapping lines +above it in `__call__` — the remap defines position semantics (comments may +disagree; trust the code). Fix preference: +1. Mode/logical permutation in the bridge spec (zero-copy) if the required + dim is already stride-1 in row-major. +2. Physical relayout only if no permutation satisfies it. + +## Symptom: `'NoneType' object is not subscriptable` during compile/trace + +You passed `None` for a parameter that is optional in another arch's +signature but required (indexed) in this one. Build the real tensor with the +shape the kernel's indexing implies, cross-checked against the vendor call +site. + +## Symptom: TypeError binding args at compile + +Argument count/order mismatch with `__call__` — re-dump the signature for +*this* arch's class (they differ between arch variants, including scalar +placement). Also check: bridge dropped aliased outputs from your launcher's +expected arguments (read installed bridge source). + +## Symptom: "not a TVM-FFI tensor" (or similar wrapper-type error) + +The direct path needs an enable flag on the DSL's `from_dlpack` in this +version. Read the error; it names the flag. + +## Symptom: works on machine A, fails on machine B + +1. Fingerprint both (inspect_environment.py); diff versions — rolling + container tags and arch-lagged builds are the historical cause of + *apparent* platform bugs. +2. Fresh-process repro on B (session contamination). +3. Same *size*/config on both? (occupancy-dependent bugs masquerade as + machine-dependent). +4. Only after 1–3: consider genuine platform/driver issues, and test with + the direct path to remove the bridge from the equation. + +## Symptom: checksums differ across identical fresh runs + +Nondeterminism red flag: unwritten regions (empty-allocated buffers), +races, or cross-stream interference. Not "just floating point" when the same +binary/config/seed is used — investigate before proceeding. + +## Escalation: kernel bug vs your bug vs bridge bug + +The isolation ladder, cheapest first: +1. Trivial elementwise `@cute.kernel` through the bridge → tests the bridge + machinery itself. +2. Target kernel through the **direct path** with a value check → tests the + kernel + your contract. +3. Target kernel through the bridge → tests the integration. +4. Toggle exactly one config flag per run via env vars in the repro script. + +When a real vendor bug is isolated: the repro script with PASS/FAIL env +toggles *is* the bug report body. Include the isolation matrix (what was +ruled out), the sanitizer result *with its blind-spot caveat stated*, and +exact environment fingerprints — triagers close silent-corruption bugs as +unreproducible unless told why the sanitizer is clean. diff --git a/docs/agent-skills/jax-cudnn-frontend/references/environment-discovery.md b/docs/agent-skills/jax-cudnn-frontend/references/environment-discovery.md new file mode 100644 index 000000000..fc9cdd7d4 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/references/environment-discovery.md @@ -0,0 +1,71 @@ +# Environment discovery + +Goal: an unambiguous fingerprint of the environment, sufficient to (a) select +the right kernel class for the hardware, (b) reproduce any finding, and (c) +detect when the environment shifts underneath you. + +Run `scripts/inspect_environment.py` first; it automates most of this. What +follows explains what each field is *for* and the pitfalls. + +## What to record and why + +| Field | Command | Why it matters | +|---|---|---| +| JAX version incl. local suffix | `python -c "import jax; print(jax.__version__)"` | Dev builds like `0.11.1.dev20260803+c6ab31b9bf` — the `+hash` is the only stable identifier; two containers can differ by days of API churn | +| jaxlib / cuda plugins | `pip list \| grep -Ei "jaxlib\|jax-cuda"` | Plugin CUDA major (cu12/cu13) must match the stack | +| cudnn-frontend | `pip list \| grep cudnn-frontend` | Kernel APIs restructure between minors (e.g. class-based → functional between 1.26 and 1.27) | +| cutlass-dsl | `pip list \| grep cutlass-dsl` | The JAX bridge lives here; bridge bugs are version-specific (e.g. aliasing broken in 4.6, fixed 4.7.1) | +| CUDA toolkit | `nvcc --version \| tail -1` | Compile-time toolchain for DSL kernels | +| cuDNN runtime | `pip list \| grep nvidia-cudnn-cu` + watch startup logs | Runtime-vs-compiled mismatches print E-level noise; usually benign for DSL kernels but record it | +| GPU + capability + driver | `nvidia-smi --query-gpu=name,compute_cap,driver_version --format=csv,noheader` | **Kernel availability and class selection are gated on compute capability.** | +| GPU count | same command (one line per GPU) | Multi-GPU nodes: pin with `CUDA_VISIBLE_DEVICES` during debugging | + +## Kernel presence check + +Presence is not implied by the package being installed: + +```python +import cudnn +print([a for a in dir(cudnn) if not a.startswith('_')]) # top-level surface +import cudnn. # target module +print(dir(cudnn.)) # what it exports +``` + +If the expected name is missing, do **not** conclude it doesn't exist — the +API may have moved (search the installed tree, Phase 2) or be lazily loaded +behind `__getattr__`. Conversely, a module *existing* does not mean your GPU +supports it: look for arch-suffixed directories (`sm90_*`, `sm100_*`, …) and +arch assertions in the code. + +## Container identity (the recurring trap) + +- **Rolling tags** (`ghcr.io/nvidia/jax:jax`) move daily. Two pulls days + apart are different environments; **arm64 builds can lag x86** by several + days, so an x86 node and a Grace node pulled the same hour can carry + different JAX commits. Multiple debugging sessions were wasted on this. +- Cite environments by the JAX `+commit` suffix, not the tag. +- Once an environment validates, freeze it: + `--container-save=/path/name.sqsh` (enroot/pyxis) and use the `.sqsh` path + thereafter. Docker: pin by digest. +- Env vars worth setting for work sessions: + `XLA_PYTHON_CLIENT_PREALLOCATE=false` (shared GPUs), + `TF_CPP_MIN_LOG_LEVEL=3` (suppress XLA C++ log spam — set **before** + importing jax). + +## When two machines disagree + +If the same code passes on machine A and fails on machine B, resist the two +easy stories ("broken platform", "broken container") until you have: + +1. Fingerprinted both environments with the script and diffed them. +2. Reproduced the failure in a fresh process on B (session contamination + mimics platform bugs). +3. Checked whether the failure is *size-dependent* rather than + machine-dependent (different default sizes/occupancy across your runs). + +In the source experience, an entire "preprod platform memory corruption" +narrative — garbage tensors, cross-machine flakiness, apparent container +regressions — was ultimately one caller-side buffer under-allocation. The +platform was innocent. Machines differing in *symptom* does not mean the +machine is the cause; buffer overflows land differently in different +allocator states. diff --git a/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md b/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md new file mode 100644 index 000000000..9dc7addf8 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md @@ -0,0 +1,137 @@ +# JAX integration + +Goal: expose the kernel to JAX as a jittable operation with correct buffer +semantics — after (never before) the standalone proof in Phase 5. + +## Discover the mechanism, don't assume it + +The bridge surface changes between versions. First introspect what exists: + +```python +import cutlass.jax as cjax +print([a for a in dir(cjax) if not a.startswith('_')]) +import inspect +print(inspect.signature(cjax.cutlass_call)) # or whatever exists +``` + +Then **read the installed bridge source** (it's Python, it's short, and it is +the only authoritative statement of the conventions below): + +```bash +python -c "import cutlass.jax.primitive as p; print(p.__file__)" +# read: the call-builder (primitive.py-like) and the compile/launch wrapper +# (compile.py-like). Both file names are version-specific. +``` + +Verify each of the following against that source for *your* version. All are +known to exist in some version, and several changed or were buggy in specific +versions: + +## Conventions to verify (each burned us once) + +**Launcher convention.** The bridge traces a `@cute.jit` function you supply. +Typical contract: `launcher(stream, *inputs, *outputs, **constexpr_kwargs)` +with outputs allocated by XLA from your `output_shape_dtype` declaration. +Your launcher's job is purely argument reordering into the kernel's +`__call__` order. + +**Outputs are uninitialized.** XLA hands the launcher raw buffers. Any buffer +the kernel *accumulates into or increments* (counters, workspaces) must be +zeroed. Two patterns: +- *Zero-prologue kernel* (robust across versions): a tiny `@cute.kernel` + that flattens the tensor and writes zeros, launched inside your launcher + before the main kernel, on the same stream. Prefer this. +- *Aliased zeroed input*: pass `jnp.zeros(...)` as an input and declare + `input_output_aliases`. **Version-sensitive**: the aliasing machinery was + silently broken in one bridge release (empty outputs, no error) and its + argument-list behavior is subtle — verify in the installed source whether + aliased outputs are *removed from the launcher's output arguments* (in the + studied version they were: you write through the input arg, and the + Python-level return still contains all declared outputs). + +**Per-tensor specs.** If the bridge has a spec type (layout/mode/alignment/ +static-ness per tensor), note: (a) it usually requires exactly one spec per +tensor — no broadcasting a single spec; (b) *mode*-style logical permutations +are zero-copy views — if the kernel demands "dim X contiguous at position k" +and dim X is already contiguous in your row-major array, a mode permutation +alone satisfies it with no physical relayout; (c) *layout*-style fields ask +XLA to materialize a different physical order — only needed when no +permutation of the existing layout satisfies the stride assertions. + +**Scalars/config.** Constexpr kwargs specialize the compiled kernel. Tuples +(e.g. problem shapes) must be hashable. Prefer passing plain Python numbers +and letting the DSL coerce, over constructing DSL scalar types inside the +traced launcher, unless the installed version's own examples do otherwise. + +**Streams.** The bridge injects XLA's stream into the launcher; pass it +through to the kernel. Never fabricate stream handles inside launcher code. + +**JIT.** Wrap the calling function in `jax.jit` (config choices that alter +shapes go in `static_argnums`). Confirm integration by checking +`jax.make_jaxpr(...)` shows the op, and — more importantly — by bit-comparing +the jitted path's output against the direct path (below). + +## Keep the direct path alive + +Alongside the JAX-native wrapper, maintain a `cute.compile`-style direct +invocation of the same kernel instance (concrete arrays via the DSL's +`from_dlpack`; check whether a TVM-FFI enable flag is required by the +installed version — a missing flag produces an explicit "not a TVM-FFI +tensor" error). This path: + +- is what the vendor's own wrappers use — the best-tested route; +- runs eagerly on concrete arrays (cannot live under `jax.jit`); +- is your A/B oracle: same kernel + same tensors through both paths must be + bit-identical. Direct-correct + bridge-wrong = bridge bug (report it); + both-wrong = kernel or contract problem; both-correct = your remaining + bugs are above this layer. + +## Pointers and ownership + +- JAX↔anything transfers via DLPack are zero-copy views. If a returned array + borrows memory owned by another framework/pool, `.copy()` it into + JAX-owned memory before the owner can release it — a freed borrowed buffer + produces illegal-address crashes at *later, unrelated* operations. +- `np.asarray(jax_array)` caches: the second call returns the first call's + host copy even if device memory changed since. Force a fresh device read + with `np.asarray(jnp.copy(x))` when auditing device state. + +## Autodiff and batching + +- The custom call defines no gradient. Provide `jax.custom_vjp`: forward + returns residuals (typically the outputs the backward kernel consumes — + their shapes discovered via a full contract pass on the *backward* + kernel, which is a separate kernel with its own arch variants, workspace, + and metadata pipeline, often including auxiliary kernels such as index + inverters). +- `jax.vmap` does not flow through custom calls; map over batch dimensions + the kernel natively supports instead. +- If any stage must run eagerly (e.g. a sub-kernel broken under the bridge + in the installed version), the enclosing training step cannot be jitted — + jit the pure-JAX portions separately and document the constraint. + +## Example (version-specific) + +The following worked against nvidia-cutlass-dsl 4.6.1/4.7.0 + cudnn-frontend +1.27.0. It illustrates the *shape* of a solution, not a timeless API — every +name below must be re-verified per Phase 4. + +```python +@cute.jit +def _launcher(stream, x, meta, out, aux, *, scale: float): + _zero_f32(aux).launch( # prologue: aux is accumulated into + grid=(cute.ceil_div(cute.size(aux), 256), 1, 1), + block=(256, 1, 1), stream=stream) + _kernel(x, out, aux, meta, cutlass.Float32(scale), stream) + +@functools.partial(jax.jit, static_argnums=()) +def op(x, meta): + return cjax.cutlass_call( + _launcher, + output_shape_dtype=[ + jax.ShapeDtypeStruct(x.shape, x.dtype), # out + jax.ShapeDtypeStruct(aux_shape, jnp.float32), # aux — shape from + ], # vendor allocation! + softmax_scale=..., + )(x, meta) +``` diff --git a/docs/agent-skills/jax-cudnn-frontend/references/validation.md b/docs/agent-skills/jax-cudnn-frontend/references/validation.md new file mode 100644 index 000000000..c436082f6 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/references/validation.md @@ -0,0 +1,106 @@ +# Validation + +Correctness gates. A kernel integration is not done until every applicable +gate passes. The source experience contains multiple runs that "passed" by +weaker standards and were later shown to be silently wrong. + +## Gate 0 — Understand what does NOT count as evidence + +- Printed shapes/dtypes: available before the kernel even executes (async). +- `block_until_ready()` returning: syncs a stream, does not check values, and + may not even surface faults from work on other streams. +- Clean `compute-sanitizer --tool memcheck`: blind to (a) writes that land in + *valid neighboring allocations* (buffer overflow between live tensors) and + (b) output regions the kernel *failed to write*. +- Loss "going down" alone: verify gradient values first; a training loop can + descend on partially-wrong gradients. + +## Gate 1 — Independent mathematical reference + +Implement the operation independently (dense/masked formulation, or a +gather-based one) and compare element-wise. + +- **Two independent references beat one.** Cross-validate them against each + other on CPU first (they agreed to ~5e-7 in the source work, which then + exonerated the reference during a kernel-vs-reference dispute). +- **Run the reference on the host (numpy)** when the GPU is under suspicion — + a GPU-side reference shares allocators, streams, and library bugs with the + thing you're testing. +- Numerical hygiene: mask with a large finite value (`-1e30`), not `-inf` + (softmax-of-`-inf` NaN semantics differ across versions/backends). +- Tolerances: bf16 inputs with fp32 accumulation → expect max-abs error + around 1e-3–1e-2 relative to output scale. Compare in fp32. An error of + ~0.3 on O(1) outputs is not "loose tolerance", it is *wrong-values* — + typically a wrong attended set / permuted metadata, not precision. + +## Gate 2 — Analytic invariants + +Cheap, exact, and independent of any reference implementation. Examples that +caught / confirmed real behavior: +- softmax rows sum to 1 ⇒ for attention backward, `sum(dV) == sum(dOut)` + exactly (up to accumulation rounding); +- uniform inputs ⇒ known closed-form outputs; +- conservation-style checksums that must match across configurations. +Assert at least one such invariant in the shipped artifact, not just during +development. + +## Gate 3 — Coverage diagnostic + +Allocate outputs with `jnp.empty` during testing so unwritten regions surface +as NaN garbage, then map defects to work units: + +```python +bad = jnp.isnan(out.astype(jnp.float32)) +per_tile = bad.reshape(B, H, n_tiles, tile, D).any(axis=(0, 1, 3, 4)) +``` + +The *pattern* is diagnostic: structured tail-of-work NaNs ⇒ scheduler/config +dropping tiles past an occupancy boundary; scattered ⇒ corruption; everything +⇒ kernel never ran / wrong buffers. + +## Gate 4 — Repeated execution and determinism + +- Execute the compiled op **N≥3 times in one process**; compare checksums and + fixed-position samples across iterations. First-call-only correctness is a + known real failure mode (buffer lifetime/overflow bugs). +- Run the whole script twice in **fresh processes**: identical checksums + expected for fixed seeds. Cross-run drift ⇒ unwritten memory or a race. +- For training integrations, the loop itself is a repeated-execution test — + but only if gradient values were independently verified first (Gate 1/2). + +## Gate 5 — Size scaling + +Test at least one size where total work exceeds one hardware wave +(`work_items > SM count`), one where it doesn't, and your target size. A +scheduler bug in the source experience was **invisible at 128 tiles and +dropped 40% of output at 512 tiles** — every small smoke test passed. + +## Gate 6 — Cross-path A/B + +Same kernel, same inputs, through the JAX-native bridge and the direct +`cute.compile`-style path: outputs must be bit-identical (same kernel binary, +same math). Any difference is an integration bug by construction. + +## Gate 7 — Config-space correctness (autotuning) + +If configuration knobs are exposed (tile sizes, scheduler flags, splits): +run Gate 1 or Gate 2 **per configuration**, not just per kernel. Silently +wrong configs exist and may be faster than correct ones — a latency-only +autotuner will select them. Reject or quarantine any config that fails +values before it is timed. + +## Gate 8 — Invalid-input behavior (cheap, optional) + +Feed one deliberately wrong input (bad dtype, undersized metadata) and note +whether the stack raises or silently proceeds. If it silently proceeds, +raise your own validation into the wrapper — you now know the kernel won't +protect users. + +## Session hygiene during all gates + +- Fresh kernel/process after *any* illegal-access error — later results in a + poisoned context are meaningless, including apparently clean ones. +- Never re-run measurement cells in a notebook and trust the result; + notebooks re-executing definition cells recompile kernels and historically + produced phantom failures *and* phantom successes. Plain scripts, fresh + processes, env-var toggles. diff --git a/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md b/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md new file mode 100644 index 000000000..5fdbc38d7 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md @@ -0,0 +1,94 @@ +# Worked example: Block-Sparse Attention (BSA) forward + backward + +How the skill's phases were applied to a real kernel family. **Everything in +`[EX]` brackets is example-specific to cudnn-frontend 1.27.0 + +nvidia-cutlass-dsl 4.6.1/4.7.0 and must be re-derived for any other version. +Unbracketed statements are the general rules being illustrated.** + +## Phase 1 — Environment + +Fingerprinted per environment-discovery.md. Decision made from compute +capability: `[EX: cc 9.0 → sm90_blk64 classes; cc 10.0 → sm100_blk64 +classes]`. The GPU marketing name misled once — a "GB200" node reports B200 +GPUs at cc 10.0, so the SM100 (not a hypothetical SM120) class applied. + +## Phase 2 — Rosetta stone + +Grepping the installed package for the kernel class name found +`[EX: cudnn/block_sparse_attention/_interface.py]` — the vendor's +orchestration. This single file yielded: + +- the true output shape of the auxiliary `lse` tensor — + `[EX: torch.empty((batch, num_heads, seqlen_q)) — per-TOKEN]` — directly + contradicting stale docs that described it per-block. *General rule: + vendor allocations are the buffer contract.* Missing this caused a 64× + buffer overflow whose fallout (illegal addresses on re-execution, corrupted + metadata, phantom platform bugs on three machines) consumed the majority of + all debugging time. +- exact positional call order per arch `[EX: SM90 places scale after the + sparse-metadata tensors; SM100 places it before them]`; +- the workspace formula and its zeroing semantics `[EX: per (B,H): + 2·q_r + q_r·d + 2·k_r·d fp32, accumulator tail zeroed on the flattened + view; rounding differs per arch]`; +- a flag coupling `[EX: is_persistent = use_clc_scheduler]` — which predicted + (correctly) that the decoupled combination was never vendor-tested. The + decoupled default silently dropped ~40% of tiles past one CTA wave: found + by the coverage diagnostic, confirmed by toggling one env var per run, and + filed upstream. + +## Phase 3 — Contract highlights + +- Layout: the kernel asserted `[EX: stride==1 at position 1 for Q/K/O and + position 0 for V]`. The signature *comment* described a different dim + order than the remap code enforced — the code won. Satisfied with + zero-copy mode permutations `[EX: TensorSpec(mode=(2,3,1,0)) for Q/K/O]`; + no physical relayout needed because the contiguous dim was already right. +- Optionals: `[EX: blocksparse_num_blocks_q2k]` is `Optional` on one arch, + required (indexed unconditionally) on another → trace-time + `'NoneType' object is not subscriptable`. Built the real tensor with the + shape implied by the kernel's indexing expression. +- Arch constraints from asserts: `[EX: bf16 only; head_dim=128 for blk64 + paths, forward and backward]`. + +## Phases 4–6 — Integration + +- Both invocation paths kept alive throughout. The A/B earned its keep + twice: `[EX: proving a persistent-scheduler bug was invocation-independent + (kernel bug → upstream report), and proving a CSR-builder kernel was + correct under cute.compile but silently empty under cutlass_call + (bridge aliasing bug in 4.6, confirmed by the bridge team as known and + fixed in 4.7.1)]`. +- Zero-init handled with a prologue `@cute.kernel` writing zeros inside the + launcher — chosen over `input_output_aliases` after the aliasing bug; + works on all bridge versions and keeps the whole pipeline jittable. +- Backward = its own full contract pass: a separate kernel per arch plus an + auxiliary index-inversion kernel and a workspace, wired under + `jax.custom_vjp`. Nothing about it was inferable from the forward. + +## Phase 7 — Validation that caught real defects + +- Dense-masked reference (host-side numpy when GPU-side references became + suspect) — caught wrong-values states that shape checks blessed. +- Two independent references cross-validated to 5e-7 — settled a + kernel-vs-reference dispute in the reference's favor. +- Analytic invariant `sum(dV) == sum(dOut) == 1.0` — free exact check on the + backward, now asserted in the shipped artifact. +- `jnp.empty` + NaN coverage map — turned "output has NaN" into "the + scheduler stops issuing tiles after the first wave". +- Size scaling: every defect above was invisible at the small smoke-test + size and appeared only past one hardware wave of work. +- Repeated execution (N≥3 + fresh-process reruns with checksum comparison) — + exposed the buffer-overflow class that single-shot tests blessed. + +## Final artifact shape (general pattern) + +1. A standalone repro/validation script: env-var toggles for every config + axis, per-stage checksums, NaN counts, fixed-position samples, hard + asserts. It served as bisection harness, regression test, and the body of + two upstream bug reports without modification. +2. The integration notebook/library: forward + custom_vjp backward, one + analytic invariant asserted inline, version-sensitive workarounds + commented with the version they apply to. +3. Bug reports for anything isolated to the vendor, each carrying: exact + environment fingerprint, PASS/FAIL command matrix, isolation performed, + and the sanitizer-blind-spot caveat. diff --git a/docs/agent-skills/jax-cudnn-frontend/scripts/contract_report.py b/docs/agent-skills/jax-cudnn-frontend/scripts/contract_report.py new file mode 100644 index 000000000..117d7b272 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/scripts/contract_report.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""Kernel contract report: signatures, in-body constraints, and vendor call +sites for a cuDNN Frontend / CuTe DSL kernel class or function. + +This automates the "Rosetta stone" step: the vendor's own call sites (and +especially its buffer ALLOCATIONS near those call sites) are the ground truth +for output shapes, argument order, and workspace layout. + +Usage: + python contract_report.py cudnn.block_sparse_attention.csrc.fwd.sm100_blk64.bsa_fwd_sm100.BlockSparseAttnForwardSm100Blk64 + python contract_report.py [--package cudnn] [--context 8] + +Notes: + - Pure introspection; never executes the kernel. + - Output is evidence, not interpretation: read the printed assert/remap + lines yourself — comments in signatures can contradict the code. +""" +import argparse +import importlib +import inspect +import os +import re +import sys + +CONSTRAINT_PAT = re.compile( + r"^\s*(assert\b|.*\bcheck_dim\(|.*element_type\b|.*const_expr\(|" + r".*raise\b|.*\.launch\()" +) +ALLOC_PAT = re.compile(r"(empty|zeros|ones|full|empty_like|zeros_like)\s*\(") + + +def resolve(dotted): + parts = dotted.split(".") + for i in range(len(parts), 0, -1): + modname = ".".join(parts[:i]) + try: + mod = importlib.import_module(modname) + except ImportError: + continue + obj = mod + try: + for attr in parts[i:]: + obj = getattr(obj, attr) + except AttributeError: + continue + return obj, mod + raise SystemExit(f"could not resolve {dotted!r} — check spelling and that " + f"the package is installed (see inspect_environment.py)") + + +def show_signature(obj, name): + print(f"\n-- {name} signature " + "-" * 40) + try: + print(f"{name}{inspect.signature(obj)}") + except (TypeError, ValueError): + print("") + # Raw source of the def line(s) preserves parameter comments, which often + # carry (possibly wrong!) shape hints — print them labeled as hints. + try: + src = inspect.getsource(obj) + header = [] + depth = 0 + for line in src.splitlines(): + header.append(line) + depth += line.count("(") - line.count(")") + if depth <= 0 and header: + break + print("raw def (comments are HINTS, code below is CONTRACT):") + for line in header[:40]: + print(f" {line}") + except (OSError, TypeError): + pass + + +def show_constraints(obj): + print("\n-- in-body constraints (asserts / check_dim / dtype / launches) --") + try: + src, start = inspect.getsourcelines(obj) + except (OSError, TypeError): + print("") + return + hits = 0 + for off, line in enumerate(src): + if CONSTRAINT_PAT.match(line): + print(f" L{start + off}: {line.rstrip()}") + hits += 1 + if hits >= 60: + print(" ... (truncated)") + break + if not hits: + print(" (none matched — read __call__ manually; contracts may be " + "enforced in helpers)") + n_launch = sum(1 for l in src if ".launch(" in l) + print(f" internal .launch( count: {n_launch}" + + (" <- multi-launch body: exercise integration bridges carefully" + if n_launch > 1 else "")) + + +def show_call_sites(symbol_name, package, context): + print(f"\n-- vendor call sites for '{symbol_name}' in package " + f"'{package}' --") + try: + pkg = importlib.import_module(package) + root = os.path.dirname(pkg.__file__) + except Exception as e: # noqa: BLE001 + print(f"") + return + found = 0 + for dirpath, _dirs, files in os.walk(root): + for fname in files: + if not fname.endswith(".py"): + continue + path = os.path.join(dirpath, fname) + try: + with open(path, errors="replace") as f: + lines = f.readlines() + except OSError: + continue + for i, line in enumerate(lines): + if symbol_name in line and f"class {symbol_name}" not in line \ + and f"def {symbol_name}" not in line: + found += 1 + rel = os.path.relpath(path, root) + print(f"\n {rel}:{i + 1}") + lo, hi = max(0, i - context), min(len(lines), i + context + 1) + for j in range(lo, hi): + mark = ">>" if j == i else " " + print(f" {mark} {j + 1}: {lines[j].rstrip()}") + # nearby allocations = candidate buffer contracts + allocs = [f" L{j + 1}: {lines[j].strip()}" + for j in range(max(0, i - 60), + min(len(lines), i + 60)) + if ALLOC_PAT.search(lines[j])] + if allocs: + print(" nearby allocations (BUFFER-SHAPE GROUND TRUTH" + " candidates):") + for a in allocs[:12]: + print(a) + if found >= 8: + print("\n ... (more call sites exist; refine manually)") + return + if not found: + print(" none found — the kernel may be exercised only from tests or " + "another package; widen the search (grep site-packages).") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("dotted", help="dotted path to kernel class/function") + ap.add_argument("--package", default=None, + help="package to scan for call sites (default: top-level " + "of the dotted path)") + ap.add_argument("--context", type=int, default=8) + args = ap.parse_args() + + obj, mod = resolve(args.dotted) + name = args.dotted.rsplit(".", 1)[-1] + print(f"resolved : {obj!r}") + print(f"defined : {getattr(mod, '__file__', '?')}") + + if inspect.isclass(obj): + show_signature(obj.__init__, f"{name}.__init__") + call = getattr(obj, "__call__", None) + if call is not None and call is not object.__call__: + show_signature(call, f"{name}.__call__") + show_constraints(call) + else: + show_signature(obj, name) + show_constraints(obj) + + package = args.package or args.dotted.split(".")[0] + show_call_sites(name, package, args.context) + + print("\n-- next steps (see references/contract-discovery.md) --") + print(" 1. Fill the contract checklist; every buffer shape must come from") + print(" a vendor allocation, not a comment or a name.") + print(" 2. Where a signature comment and an assert/remap disagree, the") + print(" code wins; record the discrepancy.") + print(" 3. Repeat for EACH architecture variant (sm90/sm100/...): ctor") + print(" params, arg order, and optionality all may differ.") + + +if __name__ == "__main__": + main() diff --git a/docs/agent-skills/jax-cudnn-frontend/scripts/inspect_environment.py b/docs/agent-skills/jax-cudnn-frontend/scripts/inspect_environment.py new file mode 100644 index 000000000..b45341fa3 --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/scripts/inspect_environment.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Environment fingerprint for cuDNN Frontend / CuTe DSL + JAX work. + +Prints a paste-ready report: package versions, CUDA toolkit, GPU +architecture, JAX backend, the cudnn-frontend surface, and the JAX-bridge +surface. Everything is best-effort: missing components are reported, never +fatal. + +Usage: + python inspect_environment.py # full report (initializes CUDA via jax) + python inspect_environment.py --no-jax # skip jax import (no CUDA context) +""" +import importlib +import importlib.metadata as md +import os +import platform +import shutil +import subprocess +import sys + + +def sh(cmd): + try: + return subprocess.run(cmd, shell=True, capture_output=True, text=True, + timeout=30).stdout.strip() + except Exception as e: # noqa: BLE001 + return f"" + + +def pkg_version(name): + try: + return md.version(name) + except md.PackageNotFoundError: + return None + + +def section(title): + print(f"\n== {title} " + "=" * max(0, 60 - len(title))) + + +def main(): + no_jax = "--no-jax" in sys.argv + + section("host") + print(f"python : {platform.python_version()} ({sys.executable})") + print(f"platform : {platform.platform()} arch={platform.machine()}") + + section("packages (pip metadata)") + interesting = [ + "jax", "jaxlib", + "jax-cuda12-plugin", "jax-cuda12-pjrt", + "jax-cuda13-plugin", "jax-cuda13-pjrt", + "nvidia-cudnn-frontend", "nvidia-cutlass-dsl", + "nvidia-cublas", "nvidia-cublas-cu12", + "nvidia-cudnn-cu12", "nvidia-cudnn-cu13", + ] + for name in interesting: + v = pkg_version(name) + if v: + print(f"{name:24s} {v}") + missing = [n for n in ("jax", "nvidia-cudnn-frontend", "nvidia-cutlass-dsl") + if not pkg_version(n)] + if missing: + print(f"MISSING (required for this workflow): {missing}") + + section("CUDA toolkit / driver / GPU") + if shutil.which("nvcc"): + print("nvcc :", sh("nvcc --version | tail -1")) + else: + print("nvcc : not on PATH") + if shutil.which("nvidia-smi"): + out = sh("nvidia-smi --query-gpu=name,compute_cap,driver_version" + " --format=csv,noheader") + for i, line in enumerate(out.splitlines()): + print(f"GPU[{i}] : {line}") + caps = {l.split(",")[1].strip() for l in out.splitlines() if "," in l} + if caps: + major = sorted(int(c.split(".")[0]) for c in caps)[0] + print(f"compute capability major: {major} " + f"(kernel classes are usually arch-gated: look for sm{major}0_* dirs)") + else: + print("nvidia-smi: not on PATH") + + section("jax runtime") + if no_jax: + print("(skipped: --no-jax)") + else: + try: + import jax # noqa: PLC0415 + print(f"jax.__version__ : {jax.__version__}") + print(f" NOTE: the '+' suffix is the only stable identifier for") + print(f" dev/nightly builds; cite it, not container tags.") + print(f"jax.__file__ : {jax.__file__}") + try: + print(f"devices : {jax.devices()}") + except Exception as e: # noqa: BLE001 + print(f"devices : ") + except Exception as e: # noqa: BLE001 + print(f"import jax failed: {e}") + + section("cudnn-frontend surface") + try: + cudnn = importlib.import_module("cudnn") + print(f"cudnn.__file__ : {getattr(cudnn, '__file__', '?')}") + names = [a for a in dir(cudnn) if not a.startswith("_")] + print(f"top-level names : {names}") + pkg_dir = os.path.dirname(getattr(cudnn, "__file__", "") or "") + if pkg_dir: + subs = sorted(d for d in os.listdir(pkg_dir) + if os.path.isdir(os.path.join(pkg_dir, d)) + and not d.startswith("_")) + print(f"subpackages : {subs}") + # arch-gated kernel dirs are a strong signal of per-arch contracts + hits = sh(f"find {pkg_dir} -maxdepth 4 -type d -name 'sm*' | head -20") + if hits: + print("arch-gated dirs :") + for line in hits.splitlines(): + print(f" {line}") + except Exception as e: # noqa: BLE001 + print(f"import cudnn failed: {e}") + + section("JAX bridge surface (cutlass.jax)") + try: + cjax = importlib.import_module("cutlass.jax") + names = [a for a in dir(cjax) if not a.startswith("_")] + print(f"exports : {names}") + import inspect # noqa: PLC0415 + for cand in ("cutlass_call",): + fn = getattr(cjax, cand, None) + if fn is not None: + try: + print(f"{cand}{inspect.signature(fn)}") + except (TypeError, ValueError): + print(f"{cand}: ") + for mod in ("cutlass.jax.primitive", "cutlass.jax.compile", + "cutlass.jax.types"): + try: + m = importlib.import_module(mod) + print(f"{mod} -> {m.__file__}") + except Exception: # noqa: BLE001 + pass + print("READ those source files: launcher convention, output allocation,") + print("aliasing semantics, and spec types are version-specific.") + except Exception as e: # noqa: BLE001 + print(f"import cutlass.jax failed: {e}") + + section("advisories seen at import time") + print("Watch process stderr on first GPU use for:") + print(" - cuDNN runtime-vs-compiled version messages (record, usually benign)") + print(" - cuBLAS known-issue warnings (e.g. TMEM concurrency <13.2 on Blackwell)") + print(" - anything printing 'capture'/'graph' during faults (command buffers)") + + +if __name__ == "__main__": + main() From f0b0ed37d8021adf296aa942615f04080bdfb660 Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Mon, 17 Aug 2026 14:01:06 -0700 Subject: [PATCH 2/7] Add skill README --- .../agent-skills/jax-cudnn-frontend/README.md | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 docs/agent-skills/jax-cudnn-frontend/README.md diff --git a/docs/agent-skills/jax-cudnn-frontend/README.md b/docs/agent-skills/jax-cudnn-frontend/README.md new file mode 100644 index 000000000..fe0e6689c --- /dev/null +++ b/docs/agent-skills/jax-cudnn-frontend/README.md @@ -0,0 +1,81 @@ +# jax-cudnn-frontend + +An agent skill for implementing, integrating, debugging, and validating +[cuDNN Frontend](https://github.com/NVIDIA/cudnn-frontend) / CuTe DSL kernels +from **pure JAX** — no PyTorch in the implementation path. + +The supported cudnn-frontend API surface for these kernels is PyTorch-based, +but the kernels themselves are framework-agnostic CuTe DSL classes. Calling +them from JAX means working against undocumented internals whose signatures, +tensor contracts, and supported flag combinations change between package +versions. This skill teaches an agent to **re-derive every contract from the +installed environment** instead of trusting memory, docs, or online examples — +and to validate with value checks, not just "it compiled and ran". + +## What the skill covers + +- Environment fingerprinting (versions, compute capability, arch-gated kernel + classes, container-identity pitfalls). +- Contract discovery: finding the vendor's own call site ("Rosetta stone") + inside the installed package and extracting buffer shapes, argument order, + workspace layout, layout/stride requirements, and flag couplings from it. +- JAX integration via both available paths — the `cutlass.jax.cutlass_call` + bridge and direct `cute.compile` — kept alive in parallel so failures can be + A/B-isolated to the kernel, the bridge, or the integration. +- Zero-initialization, aliasing, `custom_vjp` backward wiring, and + jit-compatible workarounds for known bridge limitations. +- Validation gates that catch defects compilation cannot: independent dense + references, analytic invariants, NaN coverage maps, size scaling past one + hardware wave, repeated-execution checksums. +- Symptom-indexed debugging decision trees for asynchronous-CUDA failure + modes (misattributed faults, silent buffer overflows, sanitizer blind + spots), and an escalation ladder for isolating kernel vs. bridge vs. + user bugs — ending in a file-able repro script. + +## Layout + +| Path | Purpose | +|---|---| +| `SKILL.md` | The control plane: non-negotiable rules and the 8-phase workflow. Agents start here. | +| `references/environment-discovery.md` | Fingerprinting the environment; container/version pitfalls. | +| `references/contract-discovery.md` | Deriving a kernel's data contract from installed source. | +| `references/jax-integration.md` | Both invocation paths, bridge conventions, zero-init, backward wiring. | +| `references/validation.md` | Value-level validation gates. | +| `references/debugging.md` | Symptom-indexed decision trees. | +| `references/worked-example.md` | The phases applied to a real kernel family (block-sparse attention), version-specific details bracketed `[EX]`. | +| `scripts/inspect_environment.py` | Paste-ready environment fingerprint report. | +| `scripts/contract_report.py` | Dumps signatures, in-body constraints, and vendor call sites for a kernel class. | + +## Requirements + +- A CUDA GPU environment with `jax`, `nvidia-cudnn-frontend` (with the CuTe + DSL kernels, i.e. the `csrc` kernel classes), and `nvidia-cutlass-dsl` + installed. The NVIDIA JAX container is the reference environment. +- No repository clones are needed — the workflow introspects installed + packages (see `references/contract-discovery.md`). + +## Usage + +Place this directory wherever your coding agent discovers skills — e.g. +`.claude/skills/` or `~/.claude/skills/` for Claude Code, or the equivalent +location for your tool (`.cursor/`, `.codex/`, ...). Agents without a skill +mechanism can simply be pointed at `SKILL.md` as instructions — it is +self-contained. Then invoke the skill explicitly or let it trigger on +matching tasks: + +> Using the jax-cudnn-frontend skill: wrap the cuDNN Frontend CuTe DSL +> block sparse attention kernels for inference and training from pure JAX, +> exposing the kernel configuration knobs so we can autotune per problem +> shape. + +The expected output of a run is (1) a standalone, env-var-toggled +repro/validation script, (2) the integration library or notebook with a +`custom_vjp` backward and inline invariant checks, and (3) bug reports for +anything isolated to the vendor. + +## Scope and caveats + +- The skill encodes *procedure*, not APIs: everything version-specific in the + worked example is bracketed `[EX]` and must be re-derived per environment. +- Kernel classes under `csrc` are experimental and unsupported; contracts + here may change without notice between cudnn-frontend releases. From 8e18d84dcf1c7f02261edb277448ad5820b12559 Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Mon, 17 Aug 2026 15:45:36 -0700 Subject: [PATCH 3/7] Add agent skills to documentation --- docs/agent-skills/README.md | 62 +++++++++++++++++++++++++++++++++++++ docs/gpu-kernels.md | 5 +++ docs/index.yml | 6 ++++ 3 files changed, 73 insertions(+) create mode 100644 docs/agent-skills/README.md diff --git a/docs/agent-skills/README.md b/docs/agent-skills/README.md new file mode 100644 index 000000000..2e1e071bf --- /dev/null +++ b/docs/agent-skills/README.md @@ -0,0 +1,62 @@ +--- +title: Agent Skills +subtitle: Procedural skills for AI coding agents working with JAX on NVIDIA GPUs. +slug: agent-skills +--- + +JAX-Toolbox ships **agent skills**: self-contained packages of instructions, +reference documentation, and scripts that teach an AI coding agent how to do a +specific kind of JAX integration work correctly. They encode *procedure* — +how to discover APIs from the installed environment, how to validate results, +which failure modes to expect — rather than API snippets that go stale between +package versions. + +Skills live under [`docs/agent-skills/`](https://github.com/NVIDIA/JAX-Toolbox/tree/main/docs/agent-skills) +in this repository. Each skill is a directory containing: + +- `SKILL.md` — the entry point: rules and a phased workflow the agent follows. +- `references/` — deeper documentation the workflow points into as needed. +- `scripts/` — plain Python helpers the agent runs (e.g. environment + fingerprinting, API introspection). +- `README.md` — a human-facing overview of what the skill does. + +## Prerequisites + +**An agentic coding tool.** Skills are written for AI agents that can read +files and execute shell commands. The format follows the `SKILL.md` +convention (a directory with a `SKILL.md` entry point), which many agentic +tools discover automatically — but nothing in these skills depends on any +particular tool: any agent can simply be told to read the skill's `SKILL.md` +and follow it. + +**The agent must run inside the target environment.** These skills are +discovery-driven: they instruct the agent to introspect installed packages, +run probe scripts, and validate results on a GPU. Run your agent on the +machine (or inside the container) where the work will execute — an agent +without access to the GPU and the installed Python environment cannot follow +the workflow. The [NVIDIA JAX containers](https://github.com/NVIDIA/JAX-Toolbox#containers) +are the reference environment; each skill's own README lists any additional +package requirements. + +## Using a skill + +If your tool supports skills, copy (or symlink) the skill directory into +wherever it discovers them — for example `.claude/skills/` (project) or +`~/.claude/skills/` (user) for Claude Code, or the equivalent under +`.cursor/`, `.codex/`, etc. Then either mention the skill by name or just +describe the task — the skill triggers on matching requests: + +> Using the jax-cudnn-frontend skill: wrap the cuDNN Frontend CuTe DSL block +> sparse attention kernels for inference and training from pure JAX, exposing +> the kernel configuration knobs so we can autotune per problem shape. + +If your tool has no skills mechanism, include an instruction like *"Read +`docs/agent-skills//SKILL.md` and follow it for this task"* in your +prompt — the skills are plain markdown and Python, with no tool-specific +dependencies. + +## Available skills + +| Skill | Use for | +|---|---| +| [`jax-cudnn-frontend`](jax-cudnn-frontend/README.md) | Implementing, integrating, debugging, and validating cuDNN Frontend / CuTe DSL kernels from pure JAX (no PyTorch) — wrapping, calling, porting, or autotuning attention variants, GEMM fusions, and experimental `csrc` kernel classes. | diff --git a/docs/gpu-kernels.md b/docs/gpu-kernels.md index 31937c231..d51e88543 100644 --- a/docs/gpu-kernels.md +++ b/docs/gpu-kernels.md @@ -9,3 +9,8 @@ Developers need customization of the JAX stack to customize their models and to NVIDIA offers a suite of tools, libraries and kernel DSLs, helping customization. - [**Writing High-Performance CuTe DSL kernels in JAX**](https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html) + +If you work with an AI coding agent, see [Agent Skills](agent-skills/README.md) — +starting with [`jax-cudnn-frontend`](agent-skills/jax-cudnn-frontend/README.md), +a skill for integrating cuDNN Frontend / CuTe DSL kernels into pure JAX +workflows. diff --git a/docs/index.yml b/docs/index.yml index ccfac3e37..6b80c3969 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -74,6 +74,12 @@ navigation: contents: - page: GPU Kernels path: gpu-kernels.md + - section: Agent Skills + contents: + - page: Overview + path: agent-skills/README.md + - page: jax-cudnn-frontend + path: agent-skills/jax-cudnn-frontend/README.md # ==================== Reference ==================== - section: Reference From 7580dd0937a616411e71e17685020765204694de Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Thu, 27 Aug 2026 14:15:20 -0700 Subject: [PATCH 4/7] Replace jnp.empty with jnp.fill for buffer initialization --- docs/agent-skills/jax-cudnn-frontend/SKILL.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/agent-skills/jax-cudnn-frontend/SKILL.md b/docs/agent-skills/jax-cudnn-frontend/SKILL.md index 6e5d49d8b..088f2cb4b 100644 --- a/docs/agent-skills/jax-cudnn-frontend/SKILL.md +++ b/docs/agent-skills/jax-cudnn-frontend/SKILL.md @@ -97,8 +97,11 @@ Before any notebook or abstraction: a single plain-Python script that runs the kernel once with fixed seeds, prints **value-based** evidence (NaN counts, checksums, fixed-position samples), and exposes every configuration choice as an env-var toggle. This script is simultaneously your repro for -bug reports and your bisection harness. Allocate outputs with `jnp.empty` so -unwritten regions show up as NaN — a free tile-coverage diagnostic. Then run +bug reports and your bisection harness. Where you own the output buffers +(direct path), prefill them with a NaN sentinel — `jnp.full(shape, jnp.nan, +dtype)` — so unwritten regions are detectable; do **not** use `jnp.empty`, +which can return arbitrary uninitialized bits that evade `jnp.isnan` +(bridge-path variant: validation.md Gate 3). Then run it **N≥3 iterations** in one process (repeated execution is where buffer and lifetime bugs hide) and **twice in fresh processes** (identical checksums; drift is a red flag). From 9740710b517a68f2dfda9a2a7a8b6f90799a82f5 Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Thu, 27 Aug 2026 14:21:42 -0700 Subject: [PATCH 5/7] Update coverage diagnostic to prefill instead of jnp.empty --- .../references/validation.md | 18 ++++++++++++++++-- .../references/worked-example.md | 6 ++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/docs/agent-skills/jax-cudnn-frontend/references/validation.md b/docs/agent-skills/jax-cudnn-frontend/references/validation.md index c436082f6..18d34b46e 100644 --- a/docs/agent-skills/jax-cudnn-frontend/references/validation.md +++ b/docs/agent-skills/jax-cudnn-frontend/references/validation.md @@ -46,10 +46,24 @@ development. ## Gate 3 — Coverage diagnostic -Allocate outputs with `jnp.empty` during testing so unwritten regions surface -as NaN garbage, then map defects to work units: +Prefill outputs with a NaN sentinel during testing so unwritten regions are +detectable, then map defects to work units. Do **not** use `jnp.empty` for +this: it may return genuinely uninitialized memory (documented from +JAX 0.11, and never guaranteed to be NaN before that), so an unwritten tile +can contain any finite value and evade `jnp.isnan`. + +- **Direct path** (you allocate the buffers): `out = jnp.full(shape, + jnp.nan, dtype)`. Integer outputs need a finite sentinel instead — e.g. + `jnp.full(shape, -1, jnp.int32)`, then check for surviving `-1`s. +- **Bridge path** (`cutlass_call`-style, where the bridge allocates outputs + uninitialized): the caller cannot prefill. In diagnostic builds, add a + sentinel-fill prologue kernel inside the launcher — the same pattern as + the zero-prologue in jax-integration.md, writing NaN instead of zero — or + run this gate on the direct path, where the coverage finding transfers. ```python +out = jnp.full((B, H, S, D), jnp.nan, jnp.float32) # sentinel prefill +# ... kernel writes into out ... bad = jnp.isnan(out.astype(jnp.float32)) per_tile = bad.reshape(B, H, n_tiles, tile, D).any(axis=(0, 1, 3, 4)) ``` diff --git a/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md b/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md index 5fdbc38d7..f5a9d6500 100644 --- a/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md +++ b/docs/agent-skills/jax-cudnn-frontend/references/worked-example.md @@ -73,8 +73,10 @@ orchestration. This single file yielded: kernel-vs-reference dispute in the reference's favor. - Analytic invariant `sum(dV) == sum(dOut) == 1.0` — free exact check on the backward, now asserted in the shipped artifact. -- `jnp.empty` + NaN coverage map — turned "output has NaN" into "the - scheduler stops issuing tiles after the first wave". +- NaN-sentinel coverage map (validation.md Gate 3) — turned "output has NaN" + into "the scheduler stops issuing tiles after the first wave". Prefill + outputs with `jnp.full(..., jnp.nan)`, not `jnp.empty` — uninitialized + memory is not guaranteed to read as NaN. - Size scaling: every defect above was invisible at the small smoke-test size and appeared only past one hardware wave of work. - Repeated execution (N≥3 + fresh-process reruns with checksum comparison) — From 2e80d5cea33d77f7b7c540304853dd7fe886876b Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Thu, 27 Aug 2026 14:56:32 -0700 Subject: [PATCH 6/7] Update direct path verification recommendation to prefer first external memory ownership to avoid UB --- .../references/jax-integration.md | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md b/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md index 9dc7addf8..8ec0fccf9 100644 --- a/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md +++ b/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md @@ -79,13 +79,42 @@ invocation of the same kernel instance (concrete arrays via the DSL's installed version — a missing flag produces an explicit "not a TVM-FFI tensor" error). This path: -- is what the vendor's own wrappers use — the best-tested route; +- uses the same invocation *machinery* as the vendor's own wrappers + (`from_dlpack` → compile → launch) — the best-tested route to the kernel; - runs eagerly on concrete arrays (cannot live under `jax.jit`); - is your A/B oracle: same kernel + same tensors through both paths must be bit-identical. Direct-correct + bridge-wrong = bridge bug (report it); both-wrong = kernel or contract problem; both-correct = your remaining bugs are above this layer. +**Memory ownership on this path.** The kernel must write into the output +and workspace buffers, and whether that is legal depends on *who owns the +memory* — not on the machinery or on DLPack. The vendor runs this machinery +over PyTorch tensors, which are mutable by contract, so its writes are +legitimate. Reading JAX arrays through DLPack views is likewise fine — only +mutation is at issue. For the kernel-writable buffers, in order of +preference: + +1. **Own them outside JAX (recommended when an external GPU allocator is + available).** Allocate outputs/workspace with e.g. CuPy + (`cupy.full(shape, nan, dtype)`) or raw `cuda-python` allocations — + mutable-by-contract memory, so nothing is undefined. JAX arrays appear + only as read-only inputs; results are compared via `cupy.asnumpy`, or + imported into JAX with an immediate copy (`jnp.copy` / `device_put`) + *after* all writes complete. Fully contract-clean at the cost of one + debug-only dependency. +2. **JAX-owned buffers, with containment (zero extra dependencies).** + Prefill fresh `jnp.full` sentinels and let the kernel write into their + DLPack views. This is externally-visible mutation of an immutable-by- + contract array — undefined behavior per the JAX docs — kept benign by + discipline: buffers must be fresh, eagerly created, and single-purpose + (never model inputs, never anything produced under `jit`, which XLA may + alias or cache); hold Python references until after stream sync; read + results with a fresh device read (`np.asarray(jnp.copy(x))` — the + caching trap below is a symptom of this same contract); treat mutated + arrays as terminal — copy values out, never feed them back into traced + computation. Acceptable for a debugging oracle; never for production. + ## Pointers and ownership - JAX↔anything transfers via DLPack are zero-copy views. If a returned array From 529aa82ff1308fa9969b80f5ab704ebf59b8d95a Mon Sep 17 00:00:00 2001 From: Jessica De Silva Date: Thu, 27 Aug 2026 15:13:52 -0700 Subject: [PATCH 7/7] Explicitly state example is schematic but instances were tested --- .../references/jax-integration.md | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md b/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md index 8ec0fccf9..e19564095 100644 --- a/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md +++ b/docs/agent-skills/jax-cudnn-frontend/references/jax-integration.md @@ -139,21 +139,28 @@ preference: in the installed version), the enclosing training step cannot be jitted — jit the pure-JAX portions separately and document the constraint. -## Example (version-specific) +## Example (schematic, version-specific) -The following worked against nvidia-cutlass-dsl 4.6.1/4.7.0 + cudnn-frontend -1.27.0. It illustrates the *shape* of a solution, not a timeless API — every -name below must be re-verified per Phase 4. +**This is a schematic, not runnable code.** It condenses the structure of an +integration that was validated end-to-end (forward + backward, value-checked +against independent references on both Hopper and Blackwell) under +nvidia-cutlass-dsl 4.6.1 + cudnn-frontend 1.27.0. `_kernel`, `_zero_f32`, +and `aux_shape` are placeholders for the kernel instance, a zero-fill +prologue `@cute.kernel`, and the auxiliary shape taken from the vendor's own +allocation. It illustrates the *shape* of a solution, not a timeless API — +every name and convention below must be re-derived per Phase 4. ```python @cute.jit def _launcher(stream, x, meta, out, aux, *, scale: float): - _zero_f32(aux).launch( # prologue: aux is accumulated into - grid=(cute.ceil_div(cute.size(aux), 256), 1, 1), - block=(256, 1, 1), stream=stream) + # launcher convention (verify in installed bridge source): + # (stream, *inputs, *outputs, **constexpr_kwargs) + _zero_f32(aux).launch(..., stream=stream) # prologue: aux is accumulated + # into; bridge outputs are + # uninitialized — zero them here _kernel(x, out, aux, meta, cutlass.Float32(scale), stream) -@functools.partial(jax.jit, static_argnums=()) +@jax.jit def op(x, meta): return cjax.cutlass_call( _launcher, @@ -161,6 +168,6 @@ def op(x, meta): jax.ShapeDtypeStruct(x.shape, x.dtype), # out jax.ShapeDtypeStruct(aux_shape, jnp.float32), # aux — shape from ], # vendor allocation! - softmax_scale=..., - )(x, meta) + scale=1.0 / math.sqrt(head_dim), # constexpr kwarg — name must match + )(x, meta) # the launcher's keyword parameter ```