{"owner":"facebookincubator","repo":"AITemplate","hasSkills":true,"totalSkillsCount":6,"totalTokensCount":11400,"categories":["anthropic-skill"],"hasMcp":false,"mcpConfig":null,"found":[".claude/skills/cutedsl_backend_integration_plan.md",".claude/skills/testing.md",".claude/skills/architecture.md",".claude/skills/migrate-gemm-to-cutedsl.md",".claude/skills/adding-operators.md",".claude/skills/add-cutedsl-backend.md"],"skills":{".claude/skills/cutedsl_backend_integration_plan.md":"# Adding a CuTeDSL Backend to AITemplate: Execution Plan & Reusable Skill\n\nThis document captures the end-to-end process of converting an AITemplate (AIT)\noperator from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS\nPython DSL) with AOT compilation. It is generalized as a reusable pattern\napplicable to any AIT operator.\n\n---\n\n## Architecture Overview\n\n```\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CUTLASS C++ Backend (Original) │\n│ │\n│ Python (AIT codegen) C++ (generated) Binary │\n│ ┌─────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │\n│ │ Jinja2 template │─────▶│ .cu │─nvcc──▶│ .obj │ │\n│ │ (FUNC_TEMPLATE) │ │ (CUTLASS C++ │ │ │ │\n│ │ │ │ template code) │ │ │ │\n│ └─────────────────┘ └──────────────────┘ └───────┬──────────┘ │\n│ │ link │\n│ ┌──────▼──────────┐ │\n│ │ test.so │ │\n│ └─────────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CuTeDSL Backend (New) │\n│ │\n│ Python (AIT codegen + CuTeDSL AOT) C++ (thin wrapper) Binary │\n│ ┌─────────────────┐ ┌──────────────┐ ┌───────────────┐ ┌────────────┐ │\n│ │ CuTeDSL kernel │ │ cute.compile()│ │.cu │ │.obj │ │\n│ │ (@cute.jit + │──▶│ export_to_c()│──▶│(struct init + │──▶│(from nvcc) │ │\n│ │ @cute.kernel) │ │ │ │ metadata load)│ │ │ │\n│ └─────────────────┘ └──────┬───────┘ └───────────────┘ └────┬───────┘ │\n│ │ │ link │\n│ ┌─────▼────────┐ ┌─────▼───────┐ │\n│ │_cutedsl.h │ │ test.so │ │\n│ │_cutedsl.o │────────────────────▶│ (+libcuda) │ │\n│ │(pre-compiled) │ └─────────────┘ │\n│ └──────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n```\n\n### Key Difference\n\n- **CUTLASS C++ backend**: Jinja2 template IS the kernel. It generates C++ source\n that instantiates CUTLASS template classes. `nvcc` compiles everything.\n- **CuTeDSL backend**: Kernel is written in Python using CuTeDSL decorators\n (`@cute.jit`, `@cute.kernel`). `cute.compile()` produces MLIR+cubin at\n codegen time, `export_to_c()` emits `.h`+`.o`. Only a thin C++ wrapper\n needs nvcc compilation.\n\n---\n\n## Step-by-Step Integration Guide\n\n### Step 1: Write the CuTeDSL Kernel\n\n**Create**: `backend/cuda//cutedsl__sm80.py`\n\nThis file contains the pure CuTeDSL kernel implementation:\n\n```python\nimport cutlass\nimport cutlass.cute as cute\nfrom cutlass.cute.nvgpu import cpasync, warp\n\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n self.param = param\n self.cta_sync_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=128)\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, mB: cute.Tensor, ..., stream: cuda.CUstream):\n \"\"\"Host-side: create SMEM layouts, copy atoms, MMA config, launch kernel.\"\"\"\n # 1. Compute swizzled SMEM layouts\n sA_layout = cute.tile_to_shape(get_smem_layout_atom(...), ...)\n # 2. Create copy atoms (cpasync for GMEM->SMEM, universal for stores)\n gmem_copy = cute.make_tiled_copy_tv(cute.make_copy_atom(cpasync.CopyG2SOp(), ...), ...)\n # 3. Create tiled MMA\n tiled_mma = cute.make_tiled_mma(warp.MmaF16BF16Op(...), ...)\n # 4. Compute grid dimensions\n grid_dim = (m_blocks, batch_size, 1)\n # 5. Launch kernel\n self.kernel(...).launch(grid=grid_dim, block=[128, 1, 1], stream=stream)\n\n @cute.kernel\n def kernel(self, ...all_params...):\n \"\"\"Device-side: actual CUDA kernel body.\"\"\"\n tidx, _, _ = cute.arch.thread_idx()\n m_block, batch, _ = cute.arch.block_idx()\n smem = cutlass.utils.SmemAllocator()\n # ... GEMM mainloop, epilogue, store ...\n```\n\n**Key patterns**:\n\n1. **Swizzled SMEM layouts** (for bank-conflict-free access):\n ```python\n def get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = (128 if bytes_per_row % 128 == 0 else ...) // dtype_byte\n swizzle_bits = 4 if smem_k_block_size == 128 else ...\n return cute.make_composed_layout(\n cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), 0,\n cute.make_ordered_layout((8, smem_k_block_size), order=(1, 0)),\n )\n ```\n\n2. **SM80 warp-level MMA** (`warp.MmaF16BF16Op`):\n ```python\n tiled_mma = cute.make_tiled_mma(\n warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)),\n (num_warps, 1, 1),\n permutation_mnk=(num_warps * 16, 16, 16),\n )\n ```\n\n3. **cp.async for GMEM->SMEM**:\n ```python\n atom_async = cute.make_copy_atom(\n cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),\n dtype, num_bits_per_copy=128)\n gmem_tiled_copy = cute.make_tiled_copy_tv(atom_async, thread_layout, value_layout)\n ```\n\n4. **SMEM->register via ldmatrix**:\n ```python\n smem_copy_atom = cute.make_copy_atom(\n warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), dtype)\n smem_tiled_copy_A = cute.make_tiled_copy_A(smem_copy_atom, tiled_mma)\n ```\n\n5. **Activation dispatch table**:\n ```python\n ACTIVATION_MAP = {\n \"Sigmoid\": lambda x: 1.0 / (1.0 + cute.math.exp(-x)),\n \"ReLu\": lambda x: cute.arch.fmax(x, 0.0),\n \"Identity\": lambda x: x,\n }\n ```\n\n### Step 2: Create the Backend Codegen Module\n\n**Create**: `backend/cuda//_cutedsl.py`\n\nThis file registers AIT backend functions with `_cutedsl` suffix:\n\n```python\nimport functools, os, jinja2\nfrom aitemplate.backend import registry\nfrom aitemplate.backend.target import Target\n\n# C++ wrapper template - constructs CuTeDSL typed tensor structs from AIT void* pointers\nCUTEDSL_WRAPPER_TEMPLATE = jinja2.Template(\"\"\"\n#include \n#include \"{{cutedsl_header}}\"\n\nstatic {{func_name}}_cutedsl_Kernel_Metadata_t g_metadata;\nstatic bool g_metadata_loaded = false;\n\nstatic void ensure_metadata_loaded() {\n if (!g_metadata_loaded) {\n cuInit(0);\n {{func_name}}_cutedsl_Kernel_Metadata_Load(&g_metadata);\n g_metadata_loaded = true;\n }\n}\n\nvoid {{func_name}}(...AIT signature...) {\n ensure_metadata_loaded();\n\n // Construct CuTeDSL typed tensor structs:\n {{func_name}}_cutedsl_Tensor_mA_t t_mA;\n t_mA.data = a_ptr;\n t_mA.dynamic_shapes[0] = batch; // first dynamic dim\n t_mA.dynamic_shapes[1] = seq_len; // second dynamic dim\n t_mA.dynamic_strides[0] = seq_len * K; // batch stride\n\n // ... same for all tensors ...\n\n CUstream cu_stream = reinterpret_cast(stream);\n {{func_name}}_cutedsl_wrapper(&g_metadata, &t_mA, ..., cu_stream);\n}\n\"\"\")\n\n@functools.lru_cache(maxsize=32)\ndef _aot_compile_cutedsl_kernel(n0, n1, ..., output_dir, func_name, arch):\n \"\"\"AOT compile kernel and export .h + .o.\"\"\"\n from my_kernel import MySm80Kernel\n from cutlass.cute.runtime import from_dlpack\n from cutlass.cute.export import export_to_c\n\n kernel = MySm80Kernel(n0=n0, n1=n1, ...)\n\n # Create representative tensors (shapes don't matter for dynamic dims)\n q_pt = torch.zeros(4, 128, n1, device=\"cuda\", dtype=torch.float16)\n\n # Mark dynamic modes with per-mode divisibility\n q_cute = from_dlpack(q_pt, assumed_align=16)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=0, stride_order=q_pt.dim_order(), divisibility=1)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=1, stride_order=q_pt.dim_order(), divisibility=8)\n\n # Compile and export\n compiled = cute.compile(kernel, q_cute, ..., cu_stream)\n export_to_c(compiled, file_path=output_dir, file_name=f\"{func_name}_cutedsl\")\n return h_path, o_path\n\n@registry.reg(\"cuda..gen_function_cutedsl\")\ndef gen_function_cutedsl(func_attrs):\n h_path, o_path = _aot_compile_cutedsl_kernel(...)\n func_attrs[\"cutedsl_obj_path\"] = o_path # For builder to pick up\n return CUTEDSL_WRAPPER_TEMPLATE.render(...)\n\n@registry.reg(\"cuda..func_decl_cutedsl\")\ndef func_decl_cutedsl(func_attrs): ...\n\n@registry.reg(\"cuda..func_call_cutedsl\")\ndef func_call_cutedsl(func_attrs, indent=\" \"): ...\n```\n\n**Critical detail**: The C++ wrapper MUST construct CuTeDSL typed tensor structs\n(`*_Tensor_mX_t`) with proper `dynamic_shapes[]` and `dynamic_strides[]` fields.\nThe `export_to_c()` API generates these structs -- they are NOT raw `void*`.\n\n### Step 3: Add Backend Dispatch in the Op Class\n\n**Modify**: `compiler/ops//.py`\n\nAdd methods to detect CuTeDSL mode and dispatch with suffix:\n\n```python\ndef _use_cutedsl(self) -> bool:\n current_target = target.Target.current()\n return current_target._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n\ndef gen_function(self) -> str:\n suffix = self._backend_suffix()\n self._attrs[\"backend_suffix\"] = suffix # For codegen.py to pick up\n func_key = f\"{target_name}.{op_name}.gen_function{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_decl(self, func_attrs=None) -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_decl{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_call(self, func_attrs=None, indent=\" \") -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_call{suffix}\"\n return registry.get(func_key)(self._attrs, indent)\n```\n\n### Step 4: Modify the Build System\n\n#### 4.1 `backend/codegen.py` -- `gen_function_src()`\n\nPass `workdir` to `func_attrs` so the AOT compiler knows where to write:\n```python\nfunc._attrs[\"workdir\"] = prefix\n```\n\nAfter writing the `.cu` wrapper, add the pre-compiled `.o` to `file_pairs`:\n```python\ncutedsl_obj = func._attrs.get(\"cutedsl_obj_path\")\nif cutedsl_obj and os.path.exists(cutedsl_obj):\n file_pairs.append((cutedsl_obj, cutedsl_obj))\n```\n\nIn `_process_src_ops()`, read `backend_suffix` for `func_decl` / `func_call` dispatch:\n```python\nbackend_suffix = func._attrs.get(\"backend_suffix\", \"\")\nf_func_decl = registry.get(f\"{target}.{op}.func_decl{backend_suffix}\")\nf_func_call = registry.get(f\"{target}.{op}.func_call{backend_suffix}\")\n```\n\n#### 4.2 `backend/builder.py` -- Handle pre-compiled `.o` files\n\nIn `build_objs()`: skip compilation for `.o` source files:\n```python\nif src.endswith(\".o\"):\n _LOGGER.info(f\"Skipping compilation for pre-compiled object: {src}\")\n continue\n```\n\nIn `build_so()`: add `-lcuda` when CuTeDSL objects are present:\n```python\nhas_cutedsl = any(obj.endswith(\"_cutedsl.o\") for obj in objs)\nextra_libs = \" -lcuda\" if has_cutedsl else \"\"\n```\n\nIn `gen_makefile()`: use `os.path.basename()` for `.o` paths (Make runs from workdir):\n```python\nif pair[0].endswith(\".o\"):\n obj_files.append(os.path.basename(pair[0])) # NOT full path\n has_cutedsl = True\n```\n\n### Step 5: Update Target Configuration\n\n**Modify**: `backend/target.py`\n\nAdd `use_cutedsl_` kwarg support to the Target class (already passed through\n`_kwargs`; no explicit changes needed unless you want validation).\n\n### Step 6: Update `__init__.py` and Examples\n\n**Modify**: `backend/cuda//__init__.py`\n\nAdd the new module import for auto-registration:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n**Modify**: Example/test scripts\n\nAdd CLI arg to enable CuTeDSL:\n```python\ntarget = Target(..., use_cutedsl_=True)\n```\n\n### Step 7: Add BUCK Dependencies\n\nAdd `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps\nin `examples/BUCK` for any target that uses the CuTeDSL backend.\n\n---\n\n## Common Pitfalls and Debug Checklist\n\n| # | Error | Root Cause | Fix |\n|---|-------|------------|-----|\n| 1 | `ModuleNotFoundError: No module named 'cuda'` | Missing `nvidia-cutlass-dsl` BUCK dep | Add `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps |\n| 2 | `ValueError: Invalid leading dimension: 2` | `LayoutEnum.from_tensor()` only accepts 2D tensors (leading_dim 0 or 1) | Hardcode `utils.LayoutEnum.ROW_MAJOR` for 3D row-major tensors |\n| 3 | `_ScaledBasis TypeError` in SM90 TMA | CuTeDSL TMA layout algebra doesn't yet handle 3D tensors well | Use SM80 kernel (works on SM90+ hardware too) |\n| 4 | `coord and shape of view are weakly congruent: N modes vs M coords` | After `partition_S(local_tile(..., (tile, None)))`, the tile mode creates an extra dimension | Add extra `None` to indexing: `tQgQ[None, None, None, k_tile]` instead of `tQgQ[None, None, k_tile]` |\n| 5 | `ptr alignment does not meet requirement` for 128-bit cp.async | `mark_layout_dynamic()` loses alignment info | Use `mark_compact_shape_dynamic(mode=M, divisibility=8)` to preserve alignment |\n| 6 | `cute.copy expects src and dst to have static shape` | `mark_layout_dynamic()` makes ALL dims dynamic, but inner dims must be static constants | Use per-mode `mark_compact_shape_dynamic()` only on truly dynamic modes (batch, seq_len); leave inner dims (N0, N1, K) static |\n| 7 | `shape(X) of mode(0) is not divisible by divisibility(8)` | Batch dim size not divisible by the divisibility requirement | Use per-mode divisibility: batch `div=1`, M/seq_len `div=8` |\n| 8 | `argument of type \"void *\" incompatible with \"..._Tensor_mX_t *\"` | CuTeDSL `export_to_c` generates typed tensor structs, not raw pointer APIs | Construct typed tensor structs in C++ wrapper: set `.data`, `.dynamic_shapes[]`, `.dynamic_strides[]` |\n| 9 | `No rule to make target 'tmp/.../_cutedsl.o'` | Makefile has full path but Make runs from workdir | Use `os.path.basename(pair[0])` in `gen_makefile()` |\n| 10 | `DT_TEXTREL` linker warning | CuTeDSL `.o` contains text relocations | Non-fatal warning; can be ignored |\n\n---\n\n## Dynamic Tensor Marking Strategy\n\nThis is the most critical part for correctness. CuTeDSL needs to know which\ntensor dimensions are compile-time constants vs runtime-dynamic.\n\n### DO NOT use `mark_layout_dynamic()`\n\n`mark_layout_dynamic(leading_dim=N)` makes ALL shapes and strides dynamic.\nThis breaks:\n- `cute.copy` which requires static shapes for src/dst\n- cp.async alignment checks (loses divisibility info)\n- Optimal code generation (no compile-time constants)\n\n### DO use `mark_compact_shape_dynamic()` per mode\n\n```python\ndef make_cute_tensor(torch_tensor, dynamic_modes_div):\n \"\"\"\n dynamic_modes_div: list of (mode_index, divisibility) tuples.\n Only the specified modes become dynamic; all other dims stay static.\n \"\"\"\n ct = from_dlpack(torch_tensor, assumed_align=16)\n for mode, div in dynamic_modes_div:\n ct = ct.mark_compact_shape_dynamic(\n mode=mode,\n stride_order=torch_tensor.dim_order(),\n divisibility=div,\n )\n return ct\n```\n\n**Typical b2b_bmm pattern**:\n| Tensor | Shape | Dynamic Modes |\n|--------|-------|---------------|\n| Q | [batch, M, K] | batch (div=1), M (div=8) |\n| K | [batch, N0, K] | batch (div=1) only |\n| V | [batch, N0, N1] | batch (div=1) only |\n| Bias | [batch, M, N0] | batch (div=1), M (div=8) |\n| Out | [batch, M, N1] | batch (div=1), M (div=8) |\n\n### Why div=8 for M?\n\n128-bit cp.async copies 8 fp16 elements at once. The M dimension appears in\nstrides that feed cp.async, so CuTeDSL needs `divisibility=8` to guarantee\nalignment for 128-bit vectorized copies.\n\n---\n\n## Generated Artifact Reference\n\nAfter successful compilation, the workdir contains:\n\n| File | Size (typical) | Description |\n|------|----------------|-------------|\n| `_cutedsl.h` | ~4 KB | CuTeDSL header: typed tensor structs, `Metadata_Load/Unload`, `wrapper()` |\n| `_cutedsl.o` | ~50 KB | Pre-compiled object with embedded cubin(s) |\n| `.cu` | ~4 KB | Thin C++ wrapper: `cuInit`, metadata load, struct construction, launch |\n| `.obj` | ~16 KB | Compiled wrapper object |\n| `test.so` | ~1.4 MB | Final shared library (linked with `-lcuda`) |\n\n### Header structure (`_cutedsl.h`)\n\n```c\n// Metadata struct (holds CUlibrary handle)\ntypedef struct {\n CUlibrary module;\n} _cutedsl_Kernel_Metadata_t;\n\n// Typed tensor structs (one per kernel argument)\ntypedef struct {\n void *data;\n int32_t dynamic_shapes[N_DYNAMIC_SHAPES]; // runtime batch, seq_len, etc.\n int64_t dynamic_strides[N_DYNAMIC_STRIDES]; // runtime strides\n} _cutedsl_Tensor_mX_t;\n\n// Load/Unload/Wrapper functions\nvoid _cutedsl_Kernel_Metadata_Load(_cutedsl_Kernel_Metadata_t *metadata);\nvoid _cutedsl_Kernel_Metadata_Unload(_cutedsl_Kernel_Metadata_t *metadata);\nint32_t _cutedsl_wrapper(_cutedsl_Kernel_Metadata_t *metadata,\n _cutedsl_Tensor_mX_t *mX, ..., CUstream stream);\n```\n\n---\n\n## File Checklist (for classic_b2b_bmm)\n\n### Files Created\n\n| File | Path (relative to aitemplate/) |\n|------|-------------------------------|\n| SM80 kernel | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n| SM90 kernel (WIP) | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm90.py` |\n| Backend codegen | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| AOT example | `examples/07_how_to_run_pt_model/cutedsl_b2b_bmm_aot_example.py` |\n\n### Files Modified\n\n| File | Change |\n|------|--------|\n| `compiler/ops/b2b_bmm/classic_b2b_bmm.py` | Added `_use_cutedsl()`, `_backend_suffix()`, dispatch in `gen_function/decl/call` |\n| `backend/codegen.py` | Pass `workdir`, handle `.o` in `file_pairs`, read `backend_suffix` |\n| `backend/builder.py` | Skip `.o` compilation, add `-lcuda`, use `basename()` in Makefile |\n| `backend/cuda/b2b_bmm/__init__.py` | Added `classic_b2b_bmm_cutedsl` import |\n| `examples/07_how_to_run_pt_model/classic_b2b_bmm_example.py` | Added `--use-cutedsl` and `--both` CLI args |\n| `examples/BUCK` | Added `nvidia-cutlass-dsl` dependency |\n\n---\n\n## Reference Codebases\n\n| Reference | Path | What to reuse |\n|-----------|------|---------------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | SM80 warp MMA, SMEM layout, gemm/gemm_rs patterns |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` | SM90 TMA+WGMMA, pipeline patterns |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` | Complete CuTeDSL GEMM lifecycle |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` | `export_to_c()`, `dump_to_object()` API |\n| Quack PT2 compat | `fbcode/ads_mkl/ops/cute_dsl/quack/quack_rmsnorm_pt2.py` | `@functools.cache` + `cute.compile()` caching |\n| Inductor CuTeDSL | `fbcode/caffe2/torch/_inductor/codegen/cutedsl/` | Framework integration pattern |\n| Current CUTLASS backend | `backend/cuda/b2b_bmm/classic_b2b_bmm.py` | Interface contract (signature, registry keys) |\n| CUTLASS C++ kernels | `static/include/kernels/classic_b2b_bmm/` | Algorithm reference for register-passing b2b GEMM |\n\n---\n\n## Claude Skill: `add-cutedsl-backend`\n\nBelow is a reusable Claude skill specification for adding a CuTeDSL backend\nto any AIT operator.\n\n### Skill Description\n\n```\nWhen adding a CuTeDSL backend for an AITemplate operator, follow these steps:\n\n1. RESEARCH: Read the existing CUTLASS C++ backend for the target op:\n - `backend/cuda//.py` -- understand FUNC_TEMPLATE, registry keys,\n tensor shapes, tiling parameters\n - `compiler/ops//.py` -- understand gen_function/decl/call dispatch\n\n2. KERNEL: Write the CuTeDSL kernel in `backend/cuda//cutedsl__sm80.py`:\n - Class with __init__ (compile-time params), @cute.jit __call__ (host setup),\n @cute.kernel kernel (device code)\n - SM80: use warp.MmaF16BF16Op, cpasync.CopyG2SOp, warp.LdMatrix8x8x16bOp\n - Swizzled SMEM layouts via get_smem_layout_atom()\n - Named barriers for CTA sync\n\n3. CODEGEN: Create `backend/cuda//_cutedsl.py`:\n - @functools.lru_cache AOT compile function\n - C++ wrapper template with typed tensor struct construction\n - Registry functions with _cutedsl suffix\n - CRITICAL: Use mark_compact_shape_dynamic() per-mode, NOT mark_layout_dynamic()\n - CRITICAL: C++ wrapper must construct *_Tensor_mX_t structs, NOT pass void*\n\n4. DISPATCH: Modify `compiler/ops//.py`:\n - Add _use_cutedsl(), _backend_suffix()\n - Add suffix dispatch in gen_function(), gen_function_decl(), gen_function_call()\n - Store backend_suffix in self._attrs\n\n5. BUILD: Modify `backend/codegen.py` and `backend/builder.py`:\n - codegen.py: pass workdir, handle .o in file_pairs, read backend_suffix\n - builder.py: skip .o compilation, add -lcuda, use basename() for .o in Makefile\n\n6. REGISTER: Add import in `backend/cuda//__init__.py`\n\n7. TEST: Run with Target(..., use_cutedsl_=True) and validate against\n PyTorch reference (atol=1e-2 for fp16)\n\nCommon pitfalls to watch for:\n- mark_layout_dynamic makes ALL dims dynamic -> breaks static copy requirements\n- partition_S(local_tile(..., (tile, None))) adds extra modes -> count None indices\n- CuTeDSL export_to_c generates typed structs, not raw void* -> read .h header\n- Makefile .o paths must use basename (Make runs from workdir)\n- Need -lcuda for CUlibrary-based metadata loading\n```\n\n---\n\n## Execution Order for New Operators\n\n1. **Read existing backend** -- understand tensor shapes, tiling, registry keys\n2. **Write SM80 kernel** -- implement, test standalone with `cute.compile()`\n3. **AOT export test** -- verify `export_to_c()` produces valid `.h` + `.o`\n4. **Create backend codegen** -- registry functions, C++ wrapper template\n5. **Add dispatch** -- op class modifications, `__init__.py` import\n6. **Modify build system** -- codegen.py, builder.py changes\n7. **End-to-end test** -- run full AIT pipeline, compare outputs\n8. **(Optional) SM90 kernel** -- TMA + WGMMA version for Hopper\n\n---\n\n## Verification Checklist\n\n- [ ] CuTeDSL kernel compiles: `cute.compile()` succeeds\n- [ ] AOT export produces `.h` and `.o` files\n- [ ] C++ wrapper compiles with nvcc (no type mismatches)\n- [ ] `.o` links into `.so` (no undefined symbols, `-lcuda` present)\n- [ ] Runtime loads metadata: `Kernel_Metadata_Load()` succeeds\n- [ ] Numerical correctness: output matches PyTorch reference (atol<=0.02 for fp16)\n- [ ] Dynamic shapes work: different batch sizes / sequence lengths\n- [ ] Original CUTLASS backend still works (no regression)\n",".claude/skills/testing.md":"## Testing in AITemplate\n\n### Running Tests\n\n#### Python tests with pytest\n```bash\n# Run a specific test file\npython -m pytest tests/unittest/test_my_op.py -v\n\n# Run a specific test case\npython -m pytest tests/unittest/test_my_op.py::TestMyOp::test_basic -v\n```\n\n#### Buck tests (Meta-internal)\n```bash\n# Run a specific Buck test target\nbuck2 test //aitemplate/AITemplate:test_target_name\n```\n\n### Test Patterns\n\nAITemplate tests typically follow this structure:\n\n```python\nimport unittest\nimport torch\nfrom aitemplate.testing import detect_target\nfrom aitemplate.compiler import compile_model\nfrom aitemplate.frontend import Tensor, nn\n\nclass TestMyFeature(unittest.TestCase):\n def test_basic(self):\n # 1. Check for available GPU target\n target = detect_target()\n\n # 2. Build AIT graph\n X = Tensor(shape=[2, 3], name=\"X\", is_input=True)\n Y = my_op()(X)\n Y._attrs[\"is_output\"] = True\n Y._attrs[\"name\"] = \"Y\"\n\n # 3. Compile\n module = compile_model(Y, target, \"./tmp\", \"test_my_op\")\n\n # 4. Run with real data\n x_pt = torch.randn(2, 3).cuda().half()\n y = torch.empty(2, 3).cuda().half()\n module.run_with_tensors([x_pt], [y])\n\n # 5. Compare against PyTorch reference\n y_ref = torch_reference_impl(x_pt)\n torch.testing.assert_close(y, y_ref)\n```\n\n### Key Testing Utilities\n\n- `aitemplate.testing.detect_target()` — auto-detect CUDA or ROCm\n- `compile_model()` — compile a graph to a loadable module\n- `module.run_with_tensors()` — execute the compiled model\n- Tests typically use `float16` precision and compare with tolerance\n\n### FX2AIT Tests\n\nFor the PyTorch FX converter, tests are in `fx2ait/fx2ait/test/`:\n```python\n# These tests verify that PyTorch models convert correctly to AIT\n# and produce numerically equivalent results\n```\n",".claude/skills/architecture.md":"## AITemplate Architecture\n\n### Compilation Pipeline\n\nAITemplate transforms a high-level neural network graph into a compiled GPU binary through these stages:\n\n1. **Frontend (Python)**: Define or import a model using AIT operators (`python/aitemplate/compiler/ops/`)\n2. **Graph Optimization**: Fuse operators (horizontal, vertical, memory fusion) in `python/aitemplate/compiler/transform/`\n3. **Code Generation**: Emit C++ and CUDA/HIP source files via `python/aitemplate/backend/`\n4. **Compilation**: Invoke nvcc/hipcc to produce a `.so` shared library\n5. **Runtime**: Load the `.so` using the C++ runtime in `static/`\n\n### Key Modules\n\n- **`python/aitemplate/compiler/`**: Core compiler — IR, ops, graph transforms, type inference\n - `ops/`: Operator definitions (gemm, conv, elementwise, etc.)\n - `transform/`: Graph optimization passes (fusion, memory planning)\n - `base.py`: Base tensor and operator classes\n- **`python/aitemplate/backend/`**: Code generation backends\n - `cuda/`: NVIDIA-specific codegen using CUTLASS\n - `rocm/`: AMD-specific codegen using Composable Kernel\n - `common/`: Shared codegen utilities\n- **`python/aitemplate/frontend/`**: User-facing API for building models\n- **`python/aitemplate/utils/`**: Utility functions (shape inference, profiling, etc.)\n- **`static/`**: C++ runtime for loading and running compiled models\n\n### Operator Fusion Types\n\n- **Vertical fusion**: Fuse sequential ops (e.g., GEMM + bias + ReLU)\n- **Horizontal fusion**: Batch independent ops of the same type\n- **Memory fusion**: Eliminate intermediate allocations between fused ops\n\n### Backend Codegen\n\nEach backend (CUDA/ROCm) provides:\n- Kernel templates (Jinja2-based C++ templates)\n- Profiling to select optimal kernel configurations\n- Header generation for the compiled model API\n",".claude/skills/migrate-gemm-to-cutedsl.md":"## Migrating GEMM Operators from CUTLASS C++ to CuTeDSL\n\n### Overview\n\nThis skill documents the pattern for migrating any GEMM variant (gemm_rcr_bias, gemm_rrr, gemm_ccr, etc.) from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS Python DSL), using AOT-style deployment.\n\n### Layout Reference\n\nCUTLASS GEMM naming convention: `gemm_` where R=Row-major, C=Column-major.\n\n| Layout | A shape/order | B shape/order | MMA Atom | Key property |\n|--------|--------------|--------------|----------|-------------|\n| RCR | [M,K] K-major | [N,K] K-major | TN | Both K-contiguous |\n| RRR | [M,K] K-major | [K,N] N-major | TN/NT | A=K-cont, B=N-cont |\n| CCR | [K,M] M-major | [N,K] K-major | NN | A=M-cont, B=K-cont |\n\nFor SM80 with fp16: `warp.MmaF16BF16Op(Float16, Float32, (16,8,16))` with TN layout.\n\n### Step-by-Step Migration\n\n#### 1. Identify the operation\n\nRead the CUTLASS backend file (e.g., `backend/cuda/gemm_universal/gemm_rcr_bias.py`) to understand:\n- The GEMM layout (RCR, RRR, etc.)\n- The epilogue (bias, relu, bias+relu, etc.)\n- The function signature from `common_bias.py` or `common.py`\n- Whether it supports split-K\n\n#### 2. Create the CuTeDSL kernel\n\nCreate `backend/cuda/gemm_universal/cutedsl__sm80.py`.\n\nKey components to implement:\n\n```python\nclass GemmXxxSm80Kernel:\n def __init__(self, tile_m=128, tile_n=128, tile_k=32):\n self.tile_m, self.tile_n, self.tile_k = tile_m, tile_n, tile_k\n self.num_threads = 128 # 4 warps\n\n @cute.jit\n def __call__(self, mA, mB, mC, ...):\n # 1. Set up SMEM layouts using get_smem_layout_atom()\n # 2. Define SharedStorage struct\n # 3. Create copy atoms (async for load, universal for store)\n # 4. Create tiled_mma with SM80 16x8x16 atom\n # 5. Compute grid dims and launch kernel\n\n @cute.kernel\n def kernel(self, ...):\n # 1. Allocate SMEM, partition GMEM tiles\n # 2. GEMM mainloop: K-loop with cp.async + warp MMA\n # 3. Epilogue: apply bias/activation on FP32 accumulators\n # 4. Store: Reg -> SMEM -> GMEM\n```\n\n**SMEM layout pattern** (from `cutedsl_b2b_bmm_sm80.py`):\n```python\ndef get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = cutlass.const_expr(\n 128 if bytes_per_row % 128 == 0\n else (64 if bytes_per_row % 64 == 0\n else (32 if bytes_per_row % 32 == 0 else 16))\n ) // dtype_byte\n # ... swizzle computation\n return cute.make_composed_layout(swizzle, 0, ordered_layout)\n```\n\n**GEMM mainloop pattern**:\n```python\nfor k_tile in range(k_tiles):\n cute.copy(gmem_tiled_copy, src_gmem, dst_smem)\n cute.arch.cp_async_commit_group()\n cute.arch.cp_async_wait_group(0)\n barrier.arrive_and_wait()\n\n # Load first k-block from SMEM to registers\n cute.copy(smem_tiled_copy_A, tSsA[None, None, 0], tSrA_view[None, None, 0])\n cute.copy(smem_tiled_copy_B, tSsB[None, None, 0], tSrB_view[None, None, 0])\n\n for k_block in cutlass.range_constexpr(num_k_blocks):\n k_next = (k_block + 1) % num_k_blocks\n cute.copy(smem_tiled_copy_A, tSsA[None, None, k_next], ...)\n cute.copy(smem_tiled_copy_B, tSsB[None, None, k_next], ...)\n cute.gemm(tiled_mma, acc, tSrA[None, None, k_block], tSrB[None, None, k_block], acc)\n\n barrier.arrive_and_wait()\n```\n\n**Store pattern** (Reg -> SMEM -> GMEM):\n```python\nrOut = cute.make_fragment_like(acc, dtype)\nrOut.store(acc.load().to(dtype))\nsOut = cute.make_tensor(sA.iterator, sOut_layout) # reuse SMEM\n\nsmem_copy_atom_C = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype)\nsmem_tiled_copy_C = cute.make_tiled_copy_C(smem_copy_atom_C, tiled_mma)\n# ... copy from reg to SMEM, sync, then SMEM to GMEM\n```\n\n#### 3. Handle the epilogue\n\n**1D bias (broadcast)**: Load bias into SMEM as a 2D view with stride-0 in M, use `thr_mma.partition_C()` to align with accumulator layout:\n```python\n# In __call__: create 2D broadcast view\nmBias_2d = cute.make_tensor(mBias.iterator,\n cute.make_layout((M_size, N_size), stride=(0, bias_stride)))\n\n# In kernel: tile and load\ngBias = cute.local_tile(mBias_2d, (tile_m, tile_n), (m_block, n_block))\n# ... cp.async load to sBias ...\ntCsBias = thr_mma.partition_C(sBias)\n# element-wise add in FP32\n```\n\n**2D bias**: Direct tile and load (like b2b_bmm does).\n\n**Activation**: Apply element-wise in FP32 accumulator space before store.\n\n#### 4. Create AOT example\n\nCreate `examples/07_how_to_run_pt_model/cutedsl__aot_example.py`:\n\n```python\ndef make_cute_tensor(torch_tensor):\n ct = from_dlpack(torch_tensor, assumed_align=16)\n ct = ct.mark_layout_dynamic(leading_dim=len(torch_tensor.shape) - 1)\n ct = ct.mark_compact_shape_dynamic(\n mode=len(torch_tensor.shape) - 1,\n stride_order=torch_tensor.dim_order(),\n divisibility=(128 // cutlass.Float16.width),\n )\n return ct\n\ndef jit_compile_and_validate():\n # 1. Create PyTorch reference tensors\n # 2. Compute PyTorch reference output\n # 3. Create kernel instance\n # 4. cute.compile(kernel, *cute_tensors, *scalar_args, stream)\n # 5. Execute and compare (atol=1e-2, rtol=1e-2 for fp16)\n\ndef aot_export(compiled):\n # export_to_c(compiled, file_path=output_dir, file_name=name)\n # Generates: .h + .o\n```\n\n#### 5. Add BUCK target\n\n```python\npython_binary(\n # @autodeps-skip\n name = \"cutedsl__aot_example\",\n srcs = [\"07_how_to_run_pt_model/cutedsl__aot_example.py\"],\n main_module = \"aitemplate.AITemplate.examples.07_how_to_run_pt_model.cutedsl__aot_example\",\n deps = [\n \"fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl\",\n \"//aitemplate/AITemplate/python/aitemplate:aitemplate\",\n \"//caffe2:torch\",\n ],\n)\n```\n\n### Reference Implementations\n\n| File | What it demonstrates |\n|------|---------------------|\n| `backend/cuda/gemm_universal/cutedsl_gemm_rcr_bias_sm80.py` | GEMM RCR + 1D bias, 2D grid |\n| `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` | Fused B2B GEMM, 2D bias, activation |\n| `examples/.../cutedsl_gemm_rcr_bias_aot_example.py` | AOT example for GEMM |\n| `examples/.../cutedsl_b2b_bmm_aot_example.py` | AOT example for B2B BMM |\n| `ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | Reusable SM80 helpers |\n\n### Common Pitfalls\n\n1. **Alignment**: Use `mark_compact_shape_dynamic(divisibility=8)` for fp16 to satisfy 128-bit cp.async alignment\n2. **Bias broadcast**: Create 2D view with stride-0 in broadcast dim at host level (`__call__`), not inside `@cute.kernel`\n3. **SMEM reuse**: Output staging can reuse SMEM from input tiles (e.g., `sC = cute.make_tensor(sA.iterator, sC_layout)`)\n4. **MN view**: Use `_make_acc_tensor_mn_view()` to iterate accumulator in logical (M, N) order for epilogue ops\n5. **Python slicing**: CuTe tensors don't support Python slicing (`tensor[start:]`). Use `cute.local_tile()` or `cute.make_tensor()` with offset iterators\n",".claude/skills/adding-operators.md":"## Adding New Operators to AITemplate\n\n### Overview\n\nAdding a new operator involves defining it in Python, writing backend codegen, and adding tests.\n\n### Steps\n\n#### 1. Define the Operator (Python)\n\nCreate a new file in `python/aitemplate/compiler/ops/` (or the appropriate subdirectory):\n\n```python\n# python/aitemplate/compiler/ops/common/my_new_op.py\n\nfrom aitemplate.compiler.base import Operator, Tensor\n\nclass my_new_op(Operator):\n \"\"\"Description of what the op does.\"\"\"\n\n def __init__(self):\n super().__init__()\n self._attrs[\"op\"] = \"my_new_op\"\n\n def _infer_shape(self, x: Tensor) -> List[IntVar]:\n \"\"\"Compute output shape from input shapes.\"\"\"\n return x._attrs[\"shape\"] # adjust as needed\n\n def __call__(self, x: Tensor) -> Tensor:\n self._attrs[\"inputs\"] = [x]\n self._set_depth()\n output = Tensor(shape=self._infer_shape(x), src_ops={self})\n self._attrs[\"outputs\"] = [output]\n return output\n```\n\n#### 2. Write Backend Codegen\n\nAdd codegen templates in the appropriate backend directory:\n\n- **CUDA**: `python/aitemplate/backend/cuda/`\n- **ROCm**: `python/aitemplate/backend/rocm/`\n\nEach backend file typically provides:\n- `gen_function()` — generates the C++ kernel function\n- `gen_function_decl()` — generates the function declaration\n- `gen_function_call()` — generates the call site\n\n#### 3. Register the Operator\n\nMake sure the operator is importable from the ops package by adding it to the relevant `__init__.py`.\n\n#### 4. Add Tests\n\nCreate a test file in the appropriate test directory following existing patterns. Tests typically:\n- Build a small graph using the new op\n- Compile it with `compile_model()`\n- Compare outputs against a PyTorch reference implementation\n\n### Tips\n\n- Look at similar existing operators for patterns to follow\n- Gemm ops are in `python/aitemplate/compiler/ops/gemm/`\n- Elementwise ops are in `python/aitemplate/compiler/ops/common/`\n- Use `_attrs` dict to store operator metadata\n- Shape inference is critical — get it right or compilation will fail\n",".claude/skills/add-cutedsl-backend.md":"# Skill: Add CuTeDSL Backend to an AITemplate Operator\n\n## Description\n\nThis skill guides adding a CuTeDSL (CUTLASS Python DSL) backend to an existing\nAITemplate operator, replacing the Jinja2/CUTLASS C++ code generation with\nPython-based AOT compilation via `cute.compile()` + `export_to_c()`.\n\n## When to Use\n\n- Converting an existing AIT operator from CUTLASS C++ templates to CuTeDSL\n- Adding an alternative CuTeDSL backend alongside the existing CUTLASS backend\n- Creating a new GPU kernel using CuTeDSL within the AIT framework\n\n## Prerequisites\n\n- The operator already has a working CUTLASS C++ backend\n- `nvidia-cutlass-dsl` package is available in the build system\n- Target GPU is SM80+ (Ampere or newer)\n\n## Steps\n\n### 1. RESEARCH the existing backend\n\nRead these files to understand the operator interface:\n\n```\nbackend/cuda//.py # FUNC_TEMPLATE, registry keys, tensor shapes\ncompiler/ops//.py # gen_function/decl/call dispatch\n```\n\nKey things to note:\n- Function signature (what void* params, what scalar dims)\n- Tensor shapes and which dims are static vs dynamic\n- Tiling parameters (threadblock M/N/K, warp M/N/K)\n- Registry key pattern (e.g., `cuda.classic_b2b_bmm.gen_function`)\n\n### 2. WRITE the CuTeDSL kernel\n\nCreate: `backend/cuda//cutedsl__sm80.py`\n\nStructure:\n```python\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n # Store tile sizes, compile-time constants\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, ..., stream: cuda.CUstream):\n # Host-side: SMEM layouts, copy atoms, MMA config, kernel launch\n\n @cute.kernel\n def kernel(self, ...):\n # Device-side: actual CUDA kernel body\n```\n\nKey SM80 patterns:\n- `warp.MmaF16BF16Op` for MMA atoms\n- `cpasync.CopyG2SOp()` for GMEM->SMEM async copy\n- `warp.LdMatrix8x8x16bOp` for SMEM->register loads\n- Swizzled SMEM layouts via `cute.make_composed_layout(cute.make_swizzle(...))`\n- `pipeline.NamedBarrier` for CTA synchronization\n\n### 3. CREATE the backend codegen module\n\nCreate: `backend/cuda//_cutedsl.py`\n\nMust contain:\n- `@functools.lru_cache` AOT compilation function\n- C++ wrapper template (Jinja2) that constructs typed tensor structs\n- Three registry functions with `_cutedsl` suffix:\n - `cuda..gen_function_cutedsl`\n - `cuda..func_decl_cutedsl`\n - `cuda..func_call_cutedsl`\n\nCRITICAL: The C++ wrapper must construct `*_Tensor_mX_t` typed structs\n(not pass raw void*). Read the generated `.h` header to see struct fields.\n\nCRITICAL: Use `mark_compact_shape_dynamic()` per-mode, NOT `mark_layout_dynamic()`.\nOnly mark truly dynamic dims (batch, seq_len). Keep inner dims static.\n\n### 4. ADD backend dispatch\n\nModify: `compiler/ops//.py`\n\n```python\ndef _use_cutedsl(self) -> bool:\n return target.Target.current()._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n```\n\nAdd suffix dispatch in `gen_function()`, `gen_function_decl()`, `gen_function_call()`.\nStore `backend_suffix` in `self._attrs` for codegen.py to pick up.\n\n### 5. MODIFY the build system\n\n`backend/codegen.py`:\n- Pass `workdir` to `func_attrs`\n- Add CuTeDSL `.o` to `file_pairs`\n- Read `backend_suffix` for func_decl/func_call dispatch\n\n`backend/builder.py`:\n- Skip compilation for pre-compiled `.o` files\n- Add `-lcuda` when CuTeDSL objects are present\n- Use `os.path.basename()` for `.o` paths in Makefile\n\n### 6. REGISTER the new module\n\nAdd import in `backend/cuda//__init__.py`:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n### 7. TEST\n\nRun with `Target(..., use_cutedsl_=True)`.\nValidate against PyTorch reference (atol=1e-2 for fp16).\n\n## Common Pitfalls\n\n1. `mark_layout_dynamic()` makes ALL dims dynamic -> breaks static copy reqs\n2. `partition_S(local_tile(..., (tile, None)))` adds extra modes -> count None indices\n3. `export_to_c` generates typed structs, not raw void* -> read .h header\n4. Makefile .o paths must use basename (Make runs from workdir)\n5. Need `-lcuda` for CUlibrary-based metadata loading\n6. Batch dim divisibility: use div=1 for batch, div=8 for seq_len\n7. `LayoutEnum.from_tensor()` only handles 2D -> hardcode ROW_MAJOR for 3D\n\n## References\n\n| Reference | Path |\n|-----------|------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` |\n| classic_b2b_bmm impl | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| classic_b2b_bmm SM80 | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n"},"files":{".claude/skills/cutedsl_backend_integration_plan.md":"# Adding a CuTeDSL Backend to AITemplate: Execution Plan & Reusable Skill\n\nThis document captures the end-to-end process of converting an AITemplate (AIT)\noperator from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS\nPython DSL) with AOT compilation. It is generalized as a reusable pattern\napplicable to any AIT operator.\n\n---\n\n## Architecture Overview\n\n```\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CUTLASS C++ Backend (Original) │\n│ │\n│ Python (AIT codegen) C++ (generated) Binary │\n│ ┌─────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │\n│ │ Jinja2 template │─────▶│ .cu │─nvcc──▶│ .obj │ │\n│ │ (FUNC_TEMPLATE) │ │ (CUTLASS C++ │ │ │ │\n│ │ │ │ template code) │ │ │ │\n│ └─────────────────┘ └──────────────────┘ └───────┬──────────┘ │\n│ │ link │\n│ ┌──────▼──────────┐ │\n│ │ test.so │ │\n│ └─────────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CuTeDSL Backend (New) │\n│ │\n│ Python (AIT codegen + CuTeDSL AOT) C++ (thin wrapper) Binary │\n│ ┌─────────────────┐ ┌──────────────┐ ┌───────────────┐ ┌────────────┐ │\n│ │ CuTeDSL kernel │ │ cute.compile()│ │.cu │ │.obj │ │\n│ │ (@cute.jit + │──▶│ export_to_c()│──▶│(struct init + │──▶│(from nvcc) │ │\n│ │ @cute.kernel) │ │ │ │ metadata load)│ │ │ │\n│ └─────────────────┘ └──────┬───────┘ └───────────────┘ └────┬───────┘ │\n│ │ │ link │\n│ ┌─────▼────────┐ ┌─────▼───────┐ │\n│ │_cutedsl.h │ │ test.so │ │\n│ │_cutedsl.o │────────────────────▶│ (+libcuda) │ │\n│ │(pre-compiled) │ └─────────────┘ │\n│ └──────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n```\n\n### Key Difference\n\n- **CUTLASS C++ backend**: Jinja2 template IS the kernel. It generates C++ source\n that instantiates CUTLASS template classes. `nvcc` compiles everything.\n- **CuTeDSL backend**: Kernel is written in Python using CuTeDSL decorators\n (`@cute.jit`, `@cute.kernel`). `cute.compile()` produces MLIR+cubin at\n codegen time, `export_to_c()` emits `.h`+`.o`. Only a thin C++ wrapper\n needs nvcc compilation.\n\n---\n\n## Step-by-Step Integration Guide\n\n### Step 1: Write the CuTeDSL Kernel\n\n**Create**: `backend/cuda//cutedsl__sm80.py`\n\nThis file contains the pure CuTeDSL kernel implementation:\n\n```python\nimport cutlass\nimport cutlass.cute as cute\nfrom cutlass.cute.nvgpu import cpasync, warp\n\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n self.param = param\n self.cta_sync_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=128)\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, mB: cute.Tensor, ..., stream: cuda.CUstream):\n \"\"\"Host-side: create SMEM layouts, copy atoms, MMA config, launch kernel.\"\"\"\n # 1. Compute swizzled SMEM layouts\n sA_layout = cute.tile_to_shape(get_smem_layout_atom(...), ...)\n # 2. Create copy atoms (cpasync for GMEM->SMEM, universal for stores)\n gmem_copy = cute.make_tiled_copy_tv(cute.make_copy_atom(cpasync.CopyG2SOp(), ...), ...)\n # 3. Create tiled MMA\n tiled_mma = cute.make_tiled_mma(warp.MmaF16BF16Op(...), ...)\n # 4. Compute grid dimensions\n grid_dim = (m_blocks, batch_size, 1)\n # 5. Launch kernel\n self.kernel(...).launch(grid=grid_dim, block=[128, 1, 1], stream=stream)\n\n @cute.kernel\n def kernel(self, ...all_params...):\n \"\"\"Device-side: actual CUDA kernel body.\"\"\"\n tidx, _, _ = cute.arch.thread_idx()\n m_block, batch, _ = cute.arch.block_idx()\n smem = cutlass.utils.SmemAllocator()\n # ... GEMM mainloop, epilogue, store ...\n```\n\n**Key patterns**:\n\n1. **Swizzled SMEM layouts** (for bank-conflict-free access):\n ```python\n def get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = (128 if bytes_per_row % 128 == 0 else ...) // dtype_byte\n swizzle_bits = 4 if smem_k_block_size == 128 else ...\n return cute.make_composed_layout(\n cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), 0,\n cute.make_ordered_layout((8, smem_k_block_size), order=(1, 0)),\n )\n ```\n\n2. **SM80 warp-level MMA** (`warp.MmaF16BF16Op`):\n ```python\n tiled_mma = cute.make_tiled_mma(\n warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)),\n (num_warps, 1, 1),\n permutation_mnk=(num_warps * 16, 16, 16),\n )\n ```\n\n3. **cp.async for GMEM->SMEM**:\n ```python\n atom_async = cute.make_copy_atom(\n cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),\n dtype, num_bits_per_copy=128)\n gmem_tiled_copy = cute.make_tiled_copy_tv(atom_async, thread_layout, value_layout)\n ```\n\n4. **SMEM->register via ldmatrix**:\n ```python\n smem_copy_atom = cute.make_copy_atom(\n warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), dtype)\n smem_tiled_copy_A = cute.make_tiled_copy_A(smem_copy_atom, tiled_mma)\n ```\n\n5. **Activation dispatch table**:\n ```python\n ACTIVATION_MAP = {\n \"Sigmoid\": lambda x: 1.0 / (1.0 + cute.math.exp(-x)),\n \"ReLu\": lambda x: cute.arch.fmax(x, 0.0),\n \"Identity\": lambda x: x,\n }\n ```\n\n### Step 2: Create the Backend Codegen Module\n\n**Create**: `backend/cuda//_cutedsl.py`\n\nThis file registers AIT backend functions with `_cutedsl` suffix:\n\n```python\nimport functools, os, jinja2\nfrom aitemplate.backend import registry\nfrom aitemplate.backend.target import Target\n\n# C++ wrapper template - constructs CuTeDSL typed tensor structs from AIT void* pointers\nCUTEDSL_WRAPPER_TEMPLATE = jinja2.Template(\"\"\"\n#include \n#include \"{{cutedsl_header}}\"\n\nstatic {{func_name}}_cutedsl_Kernel_Metadata_t g_metadata;\nstatic bool g_metadata_loaded = false;\n\nstatic void ensure_metadata_loaded() {\n if (!g_metadata_loaded) {\n cuInit(0);\n {{func_name}}_cutedsl_Kernel_Metadata_Load(&g_metadata);\n g_metadata_loaded = true;\n }\n}\n\nvoid {{func_name}}(...AIT signature...) {\n ensure_metadata_loaded();\n\n // Construct CuTeDSL typed tensor structs:\n {{func_name}}_cutedsl_Tensor_mA_t t_mA;\n t_mA.data = a_ptr;\n t_mA.dynamic_shapes[0] = batch; // first dynamic dim\n t_mA.dynamic_shapes[1] = seq_len; // second dynamic dim\n t_mA.dynamic_strides[0] = seq_len * K; // batch stride\n\n // ... same for all tensors ...\n\n CUstream cu_stream = reinterpret_cast(stream);\n {{func_name}}_cutedsl_wrapper(&g_metadata, &t_mA, ..., cu_stream);\n}\n\"\"\")\n\n@functools.lru_cache(maxsize=32)\ndef _aot_compile_cutedsl_kernel(n0, n1, ..., output_dir, func_name, arch):\n \"\"\"AOT compile kernel and export .h + .o.\"\"\"\n from my_kernel import MySm80Kernel\n from cutlass.cute.runtime import from_dlpack\n from cutlass.cute.export import export_to_c\n\n kernel = MySm80Kernel(n0=n0, n1=n1, ...)\n\n # Create representative tensors (shapes don't matter for dynamic dims)\n q_pt = torch.zeros(4, 128, n1, device=\"cuda\", dtype=torch.float16)\n\n # Mark dynamic modes with per-mode divisibility\n q_cute = from_dlpack(q_pt, assumed_align=16)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=0, stride_order=q_pt.dim_order(), divisibility=1)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=1, stride_order=q_pt.dim_order(), divisibility=8)\n\n # Compile and export\n compiled = cute.compile(kernel, q_cute, ..., cu_stream)\n export_to_c(compiled, file_path=output_dir, file_name=f\"{func_name}_cutedsl\")\n return h_path, o_path\n\n@registry.reg(\"cuda..gen_function_cutedsl\")\ndef gen_function_cutedsl(func_attrs):\n h_path, o_path = _aot_compile_cutedsl_kernel(...)\n func_attrs[\"cutedsl_obj_path\"] = o_path # For builder to pick up\n return CUTEDSL_WRAPPER_TEMPLATE.render(...)\n\n@registry.reg(\"cuda..func_decl_cutedsl\")\ndef func_decl_cutedsl(func_attrs): ...\n\n@registry.reg(\"cuda..func_call_cutedsl\")\ndef func_call_cutedsl(func_attrs, indent=\" \"): ...\n```\n\n**Critical detail**: The C++ wrapper MUST construct CuTeDSL typed tensor structs\n(`*_Tensor_mX_t`) with proper `dynamic_shapes[]` and `dynamic_strides[]` fields.\nThe `export_to_c()` API generates these structs -- they are NOT raw `void*`.\n\n### Step 3: Add Backend Dispatch in the Op Class\n\n**Modify**: `compiler/ops//.py`\n\nAdd methods to detect CuTeDSL mode and dispatch with suffix:\n\n```python\ndef _use_cutedsl(self) -> bool:\n current_target = target.Target.current()\n return current_target._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n\ndef gen_function(self) -> str:\n suffix = self._backend_suffix()\n self._attrs[\"backend_suffix\"] = suffix # For codegen.py to pick up\n func_key = f\"{target_name}.{op_name}.gen_function{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_decl(self, func_attrs=None) -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_decl{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_call(self, func_attrs=None, indent=\" \") -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_call{suffix}\"\n return registry.get(func_key)(self._attrs, indent)\n```\n\n### Step 4: Modify the Build System\n\n#### 4.1 `backend/codegen.py` -- `gen_function_src()`\n\nPass `workdir` to `func_attrs` so the AOT compiler knows where to write:\n```python\nfunc._attrs[\"workdir\"] = prefix\n```\n\nAfter writing the `.cu` wrapper, add the pre-compiled `.o` to `file_pairs`:\n```python\ncutedsl_obj = func._attrs.get(\"cutedsl_obj_path\")\nif cutedsl_obj and os.path.exists(cutedsl_obj):\n file_pairs.append((cutedsl_obj, cutedsl_obj))\n```\n\nIn `_process_src_ops()`, read `backend_suffix` for `func_decl` / `func_call` dispatch:\n```python\nbackend_suffix = func._attrs.get(\"backend_suffix\", \"\")\nf_func_decl = registry.get(f\"{target}.{op}.func_decl{backend_suffix}\")\nf_func_call = registry.get(f\"{target}.{op}.func_call{backend_suffix}\")\n```\n\n#### 4.2 `backend/builder.py` -- Handle pre-compiled `.o` files\n\nIn `build_objs()`: skip compilation for `.o` source files:\n```python\nif src.endswith(\".o\"):\n _LOGGER.info(f\"Skipping compilation for pre-compiled object: {src}\")\n continue\n```\n\nIn `build_so()`: add `-lcuda` when CuTeDSL objects are present:\n```python\nhas_cutedsl = any(obj.endswith(\"_cutedsl.o\") for obj in objs)\nextra_libs = \" -lcuda\" if has_cutedsl else \"\"\n```\n\nIn `gen_makefile()`: use `os.path.basename()` for `.o` paths (Make runs from workdir):\n```python\nif pair[0].endswith(\".o\"):\n obj_files.append(os.path.basename(pair[0])) # NOT full path\n has_cutedsl = True\n```\n\n### Step 5: Update Target Configuration\n\n**Modify**: `backend/target.py`\n\nAdd `use_cutedsl_` kwarg support to the Target class (already passed through\n`_kwargs`; no explicit changes needed unless you want validation).\n\n### Step 6: Update `__init__.py` and Examples\n\n**Modify**: `backend/cuda//__init__.py`\n\nAdd the new module import for auto-registration:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n**Modify**: Example/test scripts\n\nAdd CLI arg to enable CuTeDSL:\n```python\ntarget = Target(..., use_cutedsl_=True)\n```\n\n### Step 7: Add BUCK Dependencies\n\nAdd `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps\nin `examples/BUCK` for any target that uses the CuTeDSL backend.\n\n---\n\n## Common Pitfalls and Debug Checklist\n\n| # | Error | Root Cause | Fix |\n|---|-------|------------|-----|\n| 1 | `ModuleNotFoundError: No module named 'cuda'` | Missing `nvidia-cutlass-dsl` BUCK dep | Add `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps |\n| 2 | `ValueError: Invalid leading dimension: 2` | `LayoutEnum.from_tensor()` only accepts 2D tensors (leading_dim 0 or 1) | Hardcode `utils.LayoutEnum.ROW_MAJOR` for 3D row-major tensors |\n| 3 | `_ScaledBasis TypeError` in SM90 TMA | CuTeDSL TMA layout algebra doesn't yet handle 3D tensors well | Use SM80 kernel (works on SM90+ hardware too) |\n| 4 | `coord and shape of view are weakly congruent: N modes vs M coords` | After `partition_S(local_tile(..., (tile, None)))`, the tile mode creates an extra dimension | Add extra `None` to indexing: `tQgQ[None, None, None, k_tile]` instead of `tQgQ[None, None, k_tile]` |\n| 5 | `ptr alignment does not meet requirement` for 128-bit cp.async | `mark_layout_dynamic()` loses alignment info | Use `mark_compact_shape_dynamic(mode=M, divisibility=8)` to preserve alignment |\n| 6 | `cute.copy expects src and dst to have static shape` | `mark_layout_dynamic()` makes ALL dims dynamic, but inner dims must be static constants | Use per-mode `mark_compact_shape_dynamic()` only on truly dynamic modes (batch, seq_len); leave inner dims (N0, N1, K) static |\n| 7 | `shape(X) of mode(0) is not divisible by divisibility(8)` | Batch dim size not divisible by the divisibility requirement | Use per-mode divisibility: batch `div=1`, M/seq_len `div=8` |\n| 8 | `argument of type \"void *\" incompatible with \"..._Tensor_mX_t *\"` | CuTeDSL `export_to_c` generates typed tensor structs, not raw pointer APIs | Construct typed tensor structs in C++ wrapper: set `.data`, `.dynamic_shapes[]`, `.dynamic_strides[]` |\n| 9 | `No rule to make target 'tmp/.../_cutedsl.o'` | Makefile has full path but Make runs from workdir | Use `os.path.basename(pair[0])` in `gen_makefile()` |\n| 10 | `DT_TEXTREL` linker warning | CuTeDSL `.o` contains text relocations | Non-fatal warning; can be ignored |\n\n---\n\n## Dynamic Tensor Marking Strategy\n\nThis is the most critical part for correctness. CuTeDSL needs to know which\ntensor dimensions are compile-time constants vs runtime-dynamic.\n\n### DO NOT use `mark_layout_dynamic()`\n\n`mark_layout_dynamic(leading_dim=N)` makes ALL shapes and strides dynamic.\nThis breaks:\n- `cute.copy` which requires static shapes for src/dst\n- cp.async alignment checks (loses divisibility info)\n- Optimal code generation (no compile-time constants)\n\n### DO use `mark_compact_shape_dynamic()` per mode\n\n```python\ndef make_cute_tensor(torch_tensor, dynamic_modes_div):\n \"\"\"\n dynamic_modes_div: list of (mode_index, divisibility) tuples.\n Only the specified modes become dynamic; all other dims stay static.\n \"\"\"\n ct = from_dlpack(torch_tensor, assumed_align=16)\n for mode, div in dynamic_modes_div:\n ct = ct.mark_compact_shape_dynamic(\n mode=mode,\n stride_order=torch_tensor.dim_order(),\n divisibility=div,\n )\n return ct\n```\n\n**Typical b2b_bmm pattern**:\n| Tensor | Shape | Dynamic Modes |\n|--------|-------|---------------|\n| Q | [batch, M, K] | batch (div=1), M (div=8) |\n| K | [batch, N0, K] | batch (div=1) only |\n| V | [batch, N0, N1] | batch (div=1) only |\n| Bias | [batch, M, N0] | batch (div=1), M (div=8) |\n| Out | [batch, M, N1] | batch (div=1), M (div=8) |\n\n### Why div=8 for M?\n\n128-bit cp.async copies 8 fp16 elements at once. The M dimension appears in\nstrides that feed cp.async, so CuTeDSL needs `divisibility=8` to guarantee\nalignment for 128-bit vectorized copies.\n\n---\n\n## Generated Artifact Reference\n\nAfter successful compilation, the workdir contains:\n\n| File | Size (typical) | Description |\n|------|----------------|-------------|\n| `_cutedsl.h` | ~4 KB | CuTeDSL header: typed tensor structs, `Metadata_Load/Unload`, `wrapper()` |\n| `_cutedsl.o` | ~50 KB | Pre-compiled object with embedded cubin(s) |\n| `.cu` | ~4 KB | Thin C++ wrapper: `cuInit`, metadata load, struct construction, launch |\n| `.obj` | ~16 KB | Compiled wrapper object |\n| `test.so` | ~1.4 MB | Final shared library (linked with `-lcuda`) |\n\n### Header structure (`_cutedsl.h`)\n\n```c\n// Metadata struct (holds CUlibrary handle)\ntypedef struct {\n CUlibrary module;\n} _cutedsl_Kernel_Metadata_t;\n\n// Typed tensor structs (one per kernel argument)\ntypedef struct {\n void *data;\n int32_t dynamic_shapes[N_DYNAMIC_SHAPES]; // runtime batch, seq_len, etc.\n int64_t dynamic_strides[N_DYNAMIC_STRIDES]; // runtime strides\n} _cutedsl_Tensor_mX_t;\n\n// Load/Unload/Wrapper functions\nvoid _cutedsl_Kernel_Metadata_Load(_cutedsl_Kernel_Metadata_t *metadata);\nvoid _cutedsl_Kernel_Metadata_Unload(_cutedsl_Kernel_Metadata_t *metadata);\nint32_t _cutedsl_wrapper(_cutedsl_Kernel_Metadata_t *metadata,\n _cutedsl_Tensor_mX_t *mX, ..., CUstream stream);\n```\n\n---\n\n## File Checklist (for classic_b2b_bmm)\n\n### Files Created\n\n| File | Path (relative to aitemplate/) |\n|------|-------------------------------|\n| SM80 kernel | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n| SM90 kernel (WIP) | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm90.py` |\n| Backend codegen | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| AOT example | `examples/07_how_to_run_pt_model/cutedsl_b2b_bmm_aot_example.py` |\n\n### Files Modified\n\n| File | Change |\n|------|--------|\n| `compiler/ops/b2b_bmm/classic_b2b_bmm.py` | Added `_use_cutedsl()`, `_backend_suffix()`, dispatch in `gen_function/decl/call` |\n| `backend/codegen.py` | Pass `workdir`, handle `.o` in `file_pairs`, read `backend_suffix` |\n| `backend/builder.py` | Skip `.o` compilation, add `-lcuda`, use `basename()` in Makefile |\n| `backend/cuda/b2b_bmm/__init__.py` | Added `classic_b2b_bmm_cutedsl` import |\n| `examples/07_how_to_run_pt_model/classic_b2b_bmm_example.py` | Added `--use-cutedsl` and `--both` CLI args |\n| `examples/BUCK` | Added `nvidia-cutlass-dsl` dependency |\n\n---\n\n## Reference Codebases\n\n| Reference | Path | What to reuse |\n|-----------|------|---------------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | SM80 warp MMA, SMEM layout, gemm/gemm_rs patterns |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` | SM90 TMA+WGMMA, pipeline patterns |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` | Complete CuTeDSL GEMM lifecycle |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` | `export_to_c()`, `dump_to_object()` API |\n| Quack PT2 compat | `fbcode/ads_mkl/ops/cute_dsl/quack/quack_rmsnorm_pt2.py` | `@functools.cache` + `cute.compile()` caching |\n| Inductor CuTeDSL | `fbcode/caffe2/torch/_inductor/codegen/cutedsl/` | Framework integration pattern |\n| Current CUTLASS backend | `backend/cuda/b2b_bmm/classic_b2b_bmm.py` | Interface contract (signature, registry keys) |\n| CUTLASS C++ kernels | `static/include/kernels/classic_b2b_bmm/` | Algorithm reference for register-passing b2b GEMM |\n\n---\n\n## Claude Skill: `add-cutedsl-backend`\n\nBelow is a reusable Claude skill specification for adding a CuTeDSL backend\nto any AIT operator.\n\n### Skill Description\n\n```\nWhen adding a CuTeDSL backend for an AITemplate operator, follow these steps:\n\n1. RESEARCH: Read the existing CUTLASS C++ backend for the target op:\n - `backend/cuda//.py` -- understand FUNC_TEMPLATE, registry keys,\n tensor shapes, tiling parameters\n - `compiler/ops//.py` -- understand gen_function/decl/call dispatch\n\n2. KERNEL: Write the CuTeDSL kernel in `backend/cuda//cutedsl__sm80.py`:\n - Class with __init__ (compile-time params), @cute.jit __call__ (host setup),\n @cute.kernel kernel (device code)\n - SM80: use warp.MmaF16BF16Op, cpasync.CopyG2SOp, warp.LdMatrix8x8x16bOp\n - Swizzled SMEM layouts via get_smem_layout_atom()\n - Named barriers for CTA sync\n\n3. CODEGEN: Create `backend/cuda//_cutedsl.py`:\n - @functools.lru_cache AOT compile function\n - C++ wrapper template with typed tensor struct construction\n - Registry functions with _cutedsl suffix\n - CRITICAL: Use mark_compact_shape_dynamic() per-mode, NOT mark_layout_dynamic()\n - CRITICAL: C++ wrapper must construct *_Tensor_mX_t structs, NOT pass void*\n\n4. DISPATCH: Modify `compiler/ops//.py`:\n - Add _use_cutedsl(), _backend_suffix()\n - Add suffix dispatch in gen_function(), gen_function_decl(), gen_function_call()\n - Store backend_suffix in self._attrs\n\n5. BUILD: Modify `backend/codegen.py` and `backend/builder.py`:\n - codegen.py: pass workdir, handle .o in file_pairs, read backend_suffix\n - builder.py: skip .o compilation, add -lcuda, use basename() for .o in Makefile\n\n6. REGISTER: Add import in `backend/cuda//__init__.py`\n\n7. TEST: Run with Target(..., use_cutedsl_=True) and validate against\n PyTorch reference (atol=1e-2 for fp16)\n\nCommon pitfalls to watch for:\n- mark_layout_dynamic makes ALL dims dynamic -> breaks static copy requirements\n- partition_S(local_tile(..., (tile, None))) adds extra modes -> count None indices\n- CuTeDSL export_to_c generates typed structs, not raw void* -> read .h header\n- Makefile .o paths must use basename (Make runs from workdir)\n- Need -lcuda for CUlibrary-based metadata loading\n```\n\n---\n\n## Execution Order for New Operators\n\n1. **Read existing backend** -- understand tensor shapes, tiling, registry keys\n2. **Write SM80 kernel** -- implement, test standalone with `cute.compile()`\n3. **AOT export test** -- verify `export_to_c()` produces valid `.h` + `.o`\n4. **Create backend codegen** -- registry functions, C++ wrapper template\n5. **Add dispatch** -- op class modifications, `__init__.py` import\n6. **Modify build system** -- codegen.py, builder.py changes\n7. **End-to-end test** -- run full AIT pipeline, compare outputs\n8. **(Optional) SM90 kernel** -- TMA + WGMMA version for Hopper\n\n---\n\n## Verification Checklist\n\n- [ ] CuTeDSL kernel compiles: `cute.compile()` succeeds\n- [ ] AOT export produces `.h` and `.o` files\n- [ ] C++ wrapper compiles with nvcc (no type mismatches)\n- [ ] `.o` links into `.so` (no undefined symbols, `-lcuda` present)\n- [ ] Runtime loads metadata: `Kernel_Metadata_Load()` succeeds\n- [ ] Numerical correctness: output matches PyTorch reference (atol<=0.02 for fp16)\n- [ ] Dynamic shapes work: different batch sizes / sequence lengths\n- [ ] Original CUTLASS backend still works (no regression)\n",".claude/skills/testing.md":"## Testing in AITemplate\n\n### Running Tests\n\n#### Python tests with pytest\n```bash\n# Run a specific test file\npython -m pytest tests/unittest/test_my_op.py -v\n\n# Run a specific test case\npython -m pytest tests/unittest/test_my_op.py::TestMyOp::test_basic -v\n```\n\n#### Buck tests (Meta-internal)\n```bash\n# Run a specific Buck test target\nbuck2 test //aitemplate/AITemplate:test_target_name\n```\n\n### Test Patterns\n\nAITemplate tests typically follow this structure:\n\n```python\nimport unittest\nimport torch\nfrom aitemplate.testing import detect_target\nfrom aitemplate.compiler import compile_model\nfrom aitemplate.frontend import Tensor, nn\n\nclass TestMyFeature(unittest.TestCase):\n def test_basic(self):\n # 1. Check for available GPU target\n target = detect_target()\n\n # 2. Build AIT graph\n X = Tensor(shape=[2, 3], name=\"X\", is_input=True)\n Y = my_op()(X)\n Y._attrs[\"is_output\"] = True\n Y._attrs[\"name\"] = \"Y\"\n\n # 3. Compile\n module = compile_model(Y, target, \"./tmp\", \"test_my_op\")\n\n # 4. Run with real data\n x_pt = torch.randn(2, 3).cuda().half()\n y = torch.empty(2, 3).cuda().half()\n module.run_with_tensors([x_pt], [y])\n\n # 5. Compare against PyTorch reference\n y_ref = torch_reference_impl(x_pt)\n torch.testing.assert_close(y, y_ref)\n```\n\n### Key Testing Utilities\n\n- `aitemplate.testing.detect_target()` — auto-detect CUDA or ROCm\n- `compile_model()` — compile a graph to a loadable module\n- `module.run_with_tensors()` — execute the compiled model\n- Tests typically use `float16` precision and compare with tolerance\n\n### FX2AIT Tests\n\nFor the PyTorch FX converter, tests are in `fx2ait/fx2ait/test/`:\n```python\n# These tests verify that PyTorch models convert correctly to AIT\n# and produce numerically equivalent results\n```\n",".claude/skills/architecture.md":"## AITemplate Architecture\n\n### Compilation Pipeline\n\nAITemplate transforms a high-level neural network graph into a compiled GPU binary through these stages:\n\n1. **Frontend (Python)**: Define or import a model using AIT operators (`python/aitemplate/compiler/ops/`)\n2. **Graph Optimization**: Fuse operators (horizontal, vertical, memory fusion) in `python/aitemplate/compiler/transform/`\n3. **Code Generation**: Emit C++ and CUDA/HIP source files via `python/aitemplate/backend/`\n4. **Compilation**: Invoke nvcc/hipcc to produce a `.so` shared library\n5. **Runtime**: Load the `.so` using the C++ runtime in `static/`\n\n### Key Modules\n\n- **`python/aitemplate/compiler/`**: Core compiler — IR, ops, graph transforms, type inference\n - `ops/`: Operator definitions (gemm, conv, elementwise, etc.)\n - `transform/`: Graph optimization passes (fusion, memory planning)\n - `base.py`: Base tensor and operator classes\n- **`python/aitemplate/backend/`**: Code generation backends\n - `cuda/`: NVIDIA-specific codegen using CUTLASS\n - `rocm/`: AMD-specific codegen using Composable Kernel\n - `common/`: Shared codegen utilities\n- **`python/aitemplate/frontend/`**: User-facing API for building models\n- **`python/aitemplate/utils/`**: Utility functions (shape inference, profiling, etc.)\n- **`static/`**: C++ runtime for loading and running compiled models\n\n### Operator Fusion Types\n\n- **Vertical fusion**: Fuse sequential ops (e.g., GEMM + bias + ReLU)\n- **Horizontal fusion**: Batch independent ops of the same type\n- **Memory fusion**: Eliminate intermediate allocations between fused ops\n\n### Backend Codegen\n\nEach backend (CUDA/ROCm) provides:\n- Kernel templates (Jinja2-based C++ templates)\n- Profiling to select optimal kernel configurations\n- Header generation for the compiled model API\n",".claude/skills/migrate-gemm-to-cutedsl.md":"## Migrating GEMM Operators from CUTLASS C++ to CuTeDSL\n\n### Overview\n\nThis skill documents the pattern for migrating any GEMM variant (gemm_rcr_bias, gemm_rrr, gemm_ccr, etc.) from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS Python DSL), using AOT-style deployment.\n\n### Layout Reference\n\nCUTLASS GEMM naming convention: `gemm_` where R=Row-major, C=Column-major.\n\n| Layout | A shape/order | B shape/order | MMA Atom | Key property |\n|--------|--------------|--------------|----------|-------------|\n| RCR | [M,K] K-major | [N,K] K-major | TN | Both K-contiguous |\n| RRR | [M,K] K-major | [K,N] N-major | TN/NT | A=K-cont, B=N-cont |\n| CCR | [K,M] M-major | [N,K] K-major | NN | A=M-cont, B=K-cont |\n\nFor SM80 with fp16: `warp.MmaF16BF16Op(Float16, Float32, (16,8,16))` with TN layout.\n\n### Step-by-Step Migration\n\n#### 1. Identify the operation\n\nRead the CUTLASS backend file (e.g., `backend/cuda/gemm_universal/gemm_rcr_bias.py`) to understand:\n- The GEMM layout (RCR, RRR, etc.)\n- The epilogue (bias, relu, bias+relu, etc.)\n- The function signature from `common_bias.py` or `common.py`\n- Whether it supports split-K\n\n#### 2. Create the CuTeDSL kernel\n\nCreate `backend/cuda/gemm_universal/cutedsl__sm80.py`.\n\nKey components to implement:\n\n```python\nclass GemmXxxSm80Kernel:\n def __init__(self, tile_m=128, tile_n=128, tile_k=32):\n self.tile_m, self.tile_n, self.tile_k = tile_m, tile_n, tile_k\n self.num_threads = 128 # 4 warps\n\n @cute.jit\n def __call__(self, mA, mB, mC, ...):\n # 1. Set up SMEM layouts using get_smem_layout_atom()\n # 2. Define SharedStorage struct\n # 3. Create copy atoms (async for load, universal for store)\n # 4. Create tiled_mma with SM80 16x8x16 atom\n # 5. Compute grid dims and launch kernel\n\n @cute.kernel\n def kernel(self, ...):\n # 1. Allocate SMEM, partition GMEM tiles\n # 2. GEMM mainloop: K-loop with cp.async + warp MMA\n # 3. Epilogue: apply bias/activation on FP32 accumulators\n # 4. Store: Reg -> SMEM -> GMEM\n```\n\n**SMEM layout pattern** (from `cutedsl_b2b_bmm_sm80.py`):\n```python\ndef get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = cutlass.const_expr(\n 128 if bytes_per_row % 128 == 0\n else (64 if bytes_per_row % 64 == 0\n else (32 if bytes_per_row % 32 == 0 else 16))\n ) // dtype_byte\n # ... swizzle computation\n return cute.make_composed_layout(swizzle, 0, ordered_layout)\n```\n\n**GEMM mainloop pattern**:\n```python\nfor k_tile in range(k_tiles):\n cute.copy(gmem_tiled_copy, src_gmem, dst_smem)\n cute.arch.cp_async_commit_group()\n cute.arch.cp_async_wait_group(0)\n barrier.arrive_and_wait()\n\n # Load first k-block from SMEM to registers\n cute.copy(smem_tiled_copy_A, tSsA[None, None, 0], tSrA_view[None, None, 0])\n cute.copy(smem_tiled_copy_B, tSsB[None, None, 0], tSrB_view[None, None, 0])\n\n for k_block in cutlass.range_constexpr(num_k_blocks):\n k_next = (k_block + 1) % num_k_blocks\n cute.copy(smem_tiled_copy_A, tSsA[None, None, k_next], ...)\n cute.copy(smem_tiled_copy_B, tSsB[None, None, k_next], ...)\n cute.gemm(tiled_mma, acc, tSrA[None, None, k_block], tSrB[None, None, k_block], acc)\n\n barrier.arrive_and_wait()\n```\n\n**Store pattern** (Reg -> SMEM -> GMEM):\n```python\nrOut = cute.make_fragment_like(acc, dtype)\nrOut.store(acc.load().to(dtype))\nsOut = cute.make_tensor(sA.iterator, sOut_layout) # reuse SMEM\n\nsmem_copy_atom_C = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype)\nsmem_tiled_copy_C = cute.make_tiled_copy_C(smem_copy_atom_C, tiled_mma)\n# ... copy from reg to SMEM, sync, then SMEM to GMEM\n```\n\n#### 3. Handle the epilogue\n\n**1D bias (broadcast)**: Load bias into SMEM as a 2D view with stride-0 in M, use `thr_mma.partition_C()` to align with accumulator layout:\n```python\n# In __call__: create 2D broadcast view\nmBias_2d = cute.make_tensor(mBias.iterator,\n cute.make_layout((M_size, N_size), stride=(0, bias_stride)))\n\n# In kernel: tile and load\ngBias = cute.local_tile(mBias_2d, (tile_m, tile_n), (m_block, n_block))\n# ... cp.async load to sBias ...\ntCsBias = thr_mma.partition_C(sBias)\n# element-wise add in FP32\n```\n\n**2D bias**: Direct tile and load (like b2b_bmm does).\n\n**Activation**: Apply element-wise in FP32 accumulator space before store.\n\n#### 4. Create AOT example\n\nCreate `examples/07_how_to_run_pt_model/cutedsl__aot_example.py`:\n\n```python\ndef make_cute_tensor(torch_tensor):\n ct = from_dlpack(torch_tensor, assumed_align=16)\n ct = ct.mark_layout_dynamic(leading_dim=len(torch_tensor.shape) - 1)\n ct = ct.mark_compact_shape_dynamic(\n mode=len(torch_tensor.shape) - 1,\n stride_order=torch_tensor.dim_order(),\n divisibility=(128 // cutlass.Float16.width),\n )\n return ct\n\ndef jit_compile_and_validate():\n # 1. Create PyTorch reference tensors\n # 2. Compute PyTorch reference output\n # 3. Create kernel instance\n # 4. cute.compile(kernel, *cute_tensors, *scalar_args, stream)\n # 5. Execute and compare (atol=1e-2, rtol=1e-2 for fp16)\n\ndef aot_export(compiled):\n # export_to_c(compiled, file_path=output_dir, file_name=name)\n # Generates: .h + .o\n```\n\n#### 5. Add BUCK target\n\n```python\npython_binary(\n # @autodeps-skip\n name = \"cutedsl__aot_example\",\n srcs = [\"07_how_to_run_pt_model/cutedsl__aot_example.py\"],\n main_module = \"aitemplate.AITemplate.examples.07_how_to_run_pt_model.cutedsl__aot_example\",\n deps = [\n \"fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl\",\n \"//aitemplate/AITemplate/python/aitemplate:aitemplate\",\n \"//caffe2:torch\",\n ],\n)\n```\n\n### Reference Implementations\n\n| File | What it demonstrates |\n|------|---------------------|\n| `backend/cuda/gemm_universal/cutedsl_gemm_rcr_bias_sm80.py` | GEMM RCR + 1D bias, 2D grid |\n| `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` | Fused B2B GEMM, 2D bias, activation |\n| `examples/.../cutedsl_gemm_rcr_bias_aot_example.py` | AOT example for GEMM |\n| `examples/.../cutedsl_b2b_bmm_aot_example.py` | AOT example for B2B BMM |\n| `ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | Reusable SM80 helpers |\n\n### Common Pitfalls\n\n1. **Alignment**: Use `mark_compact_shape_dynamic(divisibility=8)` for fp16 to satisfy 128-bit cp.async alignment\n2. **Bias broadcast**: Create 2D view with stride-0 in broadcast dim at host level (`__call__`), not inside `@cute.kernel`\n3. **SMEM reuse**: Output staging can reuse SMEM from input tiles (e.g., `sC = cute.make_tensor(sA.iterator, sC_layout)`)\n4. **MN view**: Use `_make_acc_tensor_mn_view()` to iterate accumulator in logical (M, N) order for epilogue ops\n5. **Python slicing**: CuTe tensors don't support Python slicing (`tensor[start:]`). Use `cute.local_tile()` or `cute.make_tensor()` with offset iterators\n",".claude/skills/adding-operators.md":"## Adding New Operators to AITemplate\n\n### Overview\n\nAdding a new operator involves defining it in Python, writing backend codegen, and adding tests.\n\n### Steps\n\n#### 1. Define the Operator (Python)\n\nCreate a new file in `python/aitemplate/compiler/ops/` (or the appropriate subdirectory):\n\n```python\n# python/aitemplate/compiler/ops/common/my_new_op.py\n\nfrom aitemplate.compiler.base import Operator, Tensor\n\nclass my_new_op(Operator):\n \"\"\"Description of what the op does.\"\"\"\n\n def __init__(self):\n super().__init__()\n self._attrs[\"op\"] = \"my_new_op\"\n\n def _infer_shape(self, x: Tensor) -> List[IntVar]:\n \"\"\"Compute output shape from input shapes.\"\"\"\n return x._attrs[\"shape\"] # adjust as needed\n\n def __call__(self, x: Tensor) -> Tensor:\n self._attrs[\"inputs\"] = [x]\n self._set_depth()\n output = Tensor(shape=self._infer_shape(x), src_ops={self})\n self._attrs[\"outputs\"] = [output]\n return output\n```\n\n#### 2. Write Backend Codegen\n\nAdd codegen templates in the appropriate backend directory:\n\n- **CUDA**: `python/aitemplate/backend/cuda/`\n- **ROCm**: `python/aitemplate/backend/rocm/`\n\nEach backend file typically provides:\n- `gen_function()` — generates the C++ kernel function\n- `gen_function_decl()` — generates the function declaration\n- `gen_function_call()` — generates the call site\n\n#### 3. Register the Operator\n\nMake sure the operator is importable from the ops package by adding it to the relevant `__init__.py`.\n\n#### 4. Add Tests\n\nCreate a test file in the appropriate test directory following existing patterns. Tests typically:\n- Build a small graph using the new op\n- Compile it with `compile_model()`\n- Compare outputs against a PyTorch reference implementation\n\n### Tips\n\n- Look at similar existing operators for patterns to follow\n- Gemm ops are in `python/aitemplate/compiler/ops/gemm/`\n- Elementwise ops are in `python/aitemplate/compiler/ops/common/`\n- Use `_attrs` dict to store operator metadata\n- Shape inference is critical — get it right or compilation will fail\n",".claude/skills/add-cutedsl-backend.md":"# Skill: Add CuTeDSL Backend to an AITemplate Operator\n\n## Description\n\nThis skill guides adding a CuTeDSL (CUTLASS Python DSL) backend to an existing\nAITemplate operator, replacing the Jinja2/CUTLASS C++ code generation with\nPython-based AOT compilation via `cute.compile()` + `export_to_c()`.\n\n## When to Use\n\n- Converting an existing AIT operator from CUTLASS C++ templates to CuTeDSL\n- Adding an alternative CuTeDSL backend alongside the existing CUTLASS backend\n- Creating a new GPU kernel using CuTeDSL within the AIT framework\n\n## Prerequisites\n\n- The operator already has a working CUTLASS C++ backend\n- `nvidia-cutlass-dsl` package is available in the build system\n- Target GPU is SM80+ (Ampere or newer)\n\n## Steps\n\n### 1. RESEARCH the existing backend\n\nRead these files to understand the operator interface:\n\n```\nbackend/cuda//.py # FUNC_TEMPLATE, registry keys, tensor shapes\ncompiler/ops//.py # gen_function/decl/call dispatch\n```\n\nKey things to note:\n- Function signature (what void* params, what scalar dims)\n- Tensor shapes and which dims are static vs dynamic\n- Tiling parameters (threadblock M/N/K, warp M/N/K)\n- Registry key pattern (e.g., `cuda.classic_b2b_bmm.gen_function`)\n\n### 2. WRITE the CuTeDSL kernel\n\nCreate: `backend/cuda//cutedsl__sm80.py`\n\nStructure:\n```python\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n # Store tile sizes, compile-time constants\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, ..., stream: cuda.CUstream):\n # Host-side: SMEM layouts, copy atoms, MMA config, kernel launch\n\n @cute.kernel\n def kernel(self, ...):\n # Device-side: actual CUDA kernel body\n```\n\nKey SM80 patterns:\n- `warp.MmaF16BF16Op` for MMA atoms\n- `cpasync.CopyG2SOp()` for GMEM->SMEM async copy\n- `warp.LdMatrix8x8x16bOp` for SMEM->register loads\n- Swizzled SMEM layouts via `cute.make_composed_layout(cute.make_swizzle(...))`\n- `pipeline.NamedBarrier` for CTA synchronization\n\n### 3. CREATE the backend codegen module\n\nCreate: `backend/cuda//_cutedsl.py`\n\nMust contain:\n- `@functools.lru_cache` AOT compilation function\n- C++ wrapper template (Jinja2) that constructs typed tensor structs\n- Three registry functions with `_cutedsl` suffix:\n - `cuda..gen_function_cutedsl`\n - `cuda..func_decl_cutedsl`\n - `cuda..func_call_cutedsl`\n\nCRITICAL: The C++ wrapper must construct `*_Tensor_mX_t` typed structs\n(not pass raw void*). Read the generated `.h` header to see struct fields.\n\nCRITICAL: Use `mark_compact_shape_dynamic()` per-mode, NOT `mark_layout_dynamic()`.\nOnly mark truly dynamic dims (batch, seq_len). Keep inner dims static.\n\n### 4. ADD backend dispatch\n\nModify: `compiler/ops//.py`\n\n```python\ndef _use_cutedsl(self) -> bool:\n return target.Target.current()._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n```\n\nAdd suffix dispatch in `gen_function()`, `gen_function_decl()`, `gen_function_call()`.\nStore `backend_suffix` in `self._attrs` for codegen.py to pick up.\n\n### 5. MODIFY the build system\n\n`backend/codegen.py`:\n- Pass `workdir` to `func_attrs`\n- Add CuTeDSL `.o` to `file_pairs`\n- Read `backend_suffix` for func_decl/func_call dispatch\n\n`backend/builder.py`:\n- Skip compilation for pre-compiled `.o` files\n- Add `-lcuda` when CuTeDSL objects are present\n- Use `os.path.basename()` for `.o` paths in Makefile\n\n### 6. REGISTER the new module\n\nAdd import in `backend/cuda//__init__.py`:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n### 7. TEST\n\nRun with `Target(..., use_cutedsl_=True)`.\nValidate against PyTorch reference (atol=1e-2 for fp16).\n\n## Common Pitfalls\n\n1. `mark_layout_dynamic()` makes ALL dims dynamic -> breaks static copy reqs\n2. `partition_S(local_tile(..., (tile, None)))` adds extra modes -> count None indices\n3. `export_to_c` generates typed structs, not raw void* -> read .h header\n4. Makefile .o paths must use basename (Make runs from workdir)\n5. Need `-lcuda` for CUlibrary-based metadata loading\n6. Batch dim divisibility: use div=1 for batch, div=8 for seq_len\n7. `LayoutEnum.from_tensor()` only handles 2D -> hardcode ROW_MAJOR for 3D\n\n## References\n\n| Reference | Path |\n|-----------|------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` |\n| classic_b2b_bmm impl | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| classic_b2b_bmm SM80 | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n"},"items":[{"name":"add-cutedsl-backend.md","path":".claude/skills/add-cutedsl-backend.md","rawUrl":"https://raw.githubusercontent.com/facebookincubator/AITemplate/HEAD/.claude/skills/add-cutedsl-backend.md","title":"Skills Skill","category":"anthropic-skill","format":"markdown","content":"# Skill: Add CuTeDSL Backend to an AITemplate Operator\n\n## Description\n\nThis skill guides adding a CuTeDSL (CUTLASS Python DSL) backend to an existing\nAITemplate operator, replacing the Jinja2/CUTLASS C++ code generation with\nPython-based AOT compilation via `cute.compile()` + `export_to_c()`.\n\n## When to Use\n\n- Converting an existing AIT operator from CUTLASS C++ templates to CuTeDSL\n- Adding an alternative CuTeDSL backend alongside the existing CUTLASS backend\n- Creating a new GPU kernel using CuTeDSL within the AIT framework\n\n## Prerequisites\n\n- The operator already has a working CUTLASS C++ backend\n- `nvidia-cutlass-dsl` package is available in the build system\n- Target GPU is SM80+ (Ampere or newer)\n\n## Steps\n\n### 1. RESEARCH the existing backend\n\nRead these files to understand the operator interface:\n\n```\nbackend/cuda//.py # FUNC_TEMPLATE, registry keys, tensor shapes\ncompiler/ops//.py # gen_function/decl/call dispatch\n```\n\nKey things to note:\n- Function signature (what void* params, what scalar dims)\n- Tensor shapes and which dims are static vs dynamic\n- Tiling parameters (threadblock M/N/K, warp M/N/K)\n- Registry key pattern (e.g., `cuda.classic_b2b_bmm.gen_function`)\n\n### 2. WRITE the CuTeDSL kernel\n\nCreate: `backend/cuda//cutedsl__sm80.py`\n\nStructure:\n```python\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n # Store tile sizes, compile-time constants\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, ..., stream: cuda.CUstream):\n # Host-side: SMEM layouts, copy atoms, MMA config, kernel launch\n\n @cute.kernel\n def kernel(self, ...):\n # Device-side: actual CUDA kernel body\n```\n\nKey SM80 patterns:\n- `warp.MmaF16BF16Op` for MMA atoms\n- `cpasync.CopyG2SOp()` for GMEM->SMEM async copy\n- `warp.LdMatrix8x8x16bOp` for SMEM->register loads\n- Swizzled SMEM layouts via `cute.make_composed_layout(cute.make_swizzle(...))`\n- `pipeline.NamedBarrier` for CTA synchronization\n\n### 3. CREATE the backend codegen module\n\nCreate: `backend/cuda//_cutedsl.py`\n\nMust contain:\n- `@functools.lru_cache` AOT compilation function\n- C++ wrapper template (Jinja2) that constructs typed tensor structs\n- Three registry functions with `_cutedsl` suffix:\n - `cuda..gen_function_cutedsl`\n - `cuda..func_decl_cutedsl`\n - `cuda..func_call_cutedsl`\n\nCRITICAL: The C++ wrapper must construct `*_Tensor_mX_t` typed structs\n(not pass raw void*). Read the generated `.h` header to see struct fields.\n\nCRITICAL: Use `mark_compact_shape_dynamic()` per-mode, NOT `mark_layout_dynamic()`.\nOnly mark truly dynamic dims (batch, seq_len). Keep inner dims static.\n\n### 4. ADD backend dispatch\n\nModify: `compiler/ops//.py`\n\n```python\ndef _use_cutedsl(self) -> bool:\n return target.Target.current()._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n```\n\nAdd suffix dispatch in `gen_function()`, `gen_function_decl()`, `gen_function_call()`.\nStore `backend_suffix` in `self._attrs` for codegen.py to pick up.\n\n### 5. MODIFY the build system\n\n`backend/codegen.py`:\n- Pass `workdir` to `func_attrs`\n- Add CuTeDSL `.o` to `file_pairs`\n- Read `backend_suffix` for func_decl/func_call dispatch\n\n`backend/builder.py`:\n- Skip compilation for pre-compiled `.o` files\n- Add `-lcuda` when CuTeDSL objects are present\n- Use `os.path.basename()` for `.o` paths in Makefile\n\n### 6. REGISTER the new module\n\nAdd import in `backend/cuda//__init__.py`:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n### 7. TEST\n\nRun with `Target(..., use_cutedsl_=True)`.\nValidate against PyTorch reference (atol=1e-2 for fp16).\n\n## Common Pitfalls\n\n1. `mark_layout_dynamic()` makes ALL dims dynamic -> breaks static copy reqs\n2. `partition_S(local_tile(..., (tile, None)))` adds extra modes -> count None indices\n3. `export_to_c` generates typed structs, not raw void* -> read .h header\n4. Makefile .o paths must use basename (Make runs from workdir)\n5. Need `-lcuda` for CUlibrary-based metadata loading\n6. Batch dim divisibility: use div=1 for batch, div=8 for seq_len\n7. `LayoutEnum.from_tensor()` only handles 2D -> hardcode ROW_MAJOR for 3D\n\n## References\n\n| Reference | Path |\n|-----------|------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` |\n| classic_b2b_bmm impl | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| classic_b2b_bmm SM80 | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n","isInternal":false,"tokens":1388,"sizeBytes":4808},{"name":"adding-operators.md","path":".claude/skills/adding-operators.md","rawUrl":"https://raw.githubusercontent.com/facebookincubator/AITemplate/HEAD/.claude/skills/adding-operators.md","title":"Skills Skill","category":"anthropic-skill","format":"markdown","content":"## Adding New Operators to AITemplate\n\n### Overview\n\nAdding a new operator involves defining it in Python, writing backend codegen, and adding tests.\n\n### Steps\n\n#### 1. Define the Operator (Python)\n\nCreate a new file in `python/aitemplate/compiler/ops/` (or the appropriate subdirectory):\n\n```python\n# python/aitemplate/compiler/ops/common/my_new_op.py\n\nfrom aitemplate.compiler.base import Operator, Tensor\n\nclass my_new_op(Operator):\n \"\"\"Description of what the op does.\"\"\"\n\n def __init__(self):\n super().__init__()\n self._attrs[\"op\"] = \"my_new_op\"\n\n def _infer_shape(self, x: Tensor) -> List[IntVar]:\n \"\"\"Compute output shape from input shapes.\"\"\"\n return x._attrs[\"shape\"] # adjust as needed\n\n def __call__(self, x: Tensor) -> Tensor:\n self._attrs[\"inputs\"] = [x]\n self._set_depth()\n output = Tensor(shape=self._infer_shape(x), src_ops={self})\n self._attrs[\"outputs\"] = [output]\n return output\n```\n\n#### 2. Write Backend Codegen\n\nAdd codegen templates in the appropriate backend directory:\n\n- **CUDA**: `python/aitemplate/backend/cuda/`\n- **ROCm**: `python/aitemplate/backend/rocm/`\n\nEach backend file typically provides:\n- `gen_function()` — generates the C++ kernel function\n- `gen_function_decl()` — generates the function declaration\n- `gen_function_call()` — generates the call site\n\n#### 3. Register the Operator\n\nMake sure the operator is importable from the ops package by adding it to the relevant `__init__.py`.\n\n#### 4. Add Tests\n\nCreate a test file in the appropriate test directory following existing patterns. Tests typically:\n- Build a small graph using the new op\n- Compile it with `compile_model()`\n- Compare outputs against a PyTorch reference implementation\n\n### Tips\n\n- Look at similar existing operators for patterns to follow\n- Gemm ops are in `python/aitemplate/compiler/ops/gemm/`\n- Elementwise ops are in `python/aitemplate/compiler/ops/common/`\n- Use `_attrs` dict to store operator metadata\n- Shape inference is critical — get it right or compilation will fail\n","isInternal":false,"tokens":488,"sizeBytes":2077},{"name":"architecture.md","path":".claude/skills/architecture.md","rawUrl":"https://raw.githubusercontent.com/facebookincubator/AITemplate/HEAD/.claude/skills/architecture.md","title":"Skills Skill","category":"anthropic-skill","format":"markdown","content":"## AITemplate Architecture\n\n### Compilation Pipeline\n\nAITemplate transforms a high-level neural network graph into a compiled GPU binary through these stages:\n\n1. **Frontend (Python)**: Define or import a model using AIT operators (`python/aitemplate/compiler/ops/`)\n2. **Graph Optimization**: Fuse operators (horizontal, vertical, memory fusion) in `python/aitemplate/compiler/transform/`\n3. **Code Generation**: Emit C++ and CUDA/HIP source files via `python/aitemplate/backend/`\n4. **Compilation**: Invoke nvcc/hipcc to produce a `.so` shared library\n5. **Runtime**: Load the `.so` using the C++ runtime in `static/`\n\n### Key Modules\n\n- **`python/aitemplate/compiler/`**: Core compiler — IR, ops, graph transforms, type inference\n - `ops/`: Operator definitions (gemm, conv, elementwise, etc.)\n - `transform/`: Graph optimization passes (fusion, memory planning)\n - `base.py`: Base tensor and operator classes\n- **`python/aitemplate/backend/`**: Code generation backends\n - `cuda/`: NVIDIA-specific codegen using CUTLASS\n - `rocm/`: AMD-specific codegen using Composable Kernel\n - `common/`: Shared codegen utilities\n- **`python/aitemplate/frontend/`**: User-facing API for building models\n- **`python/aitemplate/utils/`**: Utility functions (shape inference, profiling, etc.)\n- **`static/`**: C++ runtime for loading and running compiled models\n\n### Operator Fusion Types\n\n- **Vertical fusion**: Fuse sequential ops (e.g., GEMM + bias + ReLU)\n- **Horizontal fusion**: Batch independent ops of the same type\n- **Memory fusion**: Eliminate intermediate allocations between fused ops\n\n### Backend Codegen\n\nEach backend (CUDA/ROCm) provides:\n- Kernel templates (Jinja2-based C++ templates)\n- Profiling to select optimal kernel configurations\n- Header generation for the compiled model API\n","isInternal":false,"tokens":432,"sizeBytes":1797},{"name":"cutedsl_backend_integration_plan.md","path":".claude/skills/cutedsl_backend_integration_plan.md","rawUrl":"https://raw.githubusercontent.com/facebookincubator/AITemplate/HEAD/.claude/skills/cutedsl_backend_integration_plan.md","title":"Skills Skill","category":"anthropic-skill","format":"markdown","content":"# Adding a CuTeDSL Backend to AITemplate: Execution Plan & Reusable Skill\n\nThis document captures the end-to-end process of converting an AITemplate (AIT)\noperator from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS\nPython DSL) with AOT compilation. It is generalized as a reusable pattern\napplicable to any AIT operator.\n\n---\n\n## Architecture Overview\n\n```\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CUTLASS C++ Backend (Original) │\n│ │\n│ Python (AIT codegen) C++ (generated) Binary │\n│ ┌─────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │\n│ │ Jinja2 template │─────▶│ .cu │─nvcc──▶│ .obj │ │\n│ │ (FUNC_TEMPLATE) │ │ (CUTLASS C++ │ │ │ │\n│ │ │ │ template code) │ │ │ │\n│ └─────────────────┘ └──────────────────┘ └───────┬──────────┘ │\n│ │ link │\n│ ┌──────▼──────────┐ │\n│ │ test.so │ │\n│ └─────────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CuTeDSL Backend (New) │\n│ │\n│ Python (AIT codegen + CuTeDSL AOT) C++ (thin wrapper) Binary │\n│ ┌─────────────────┐ ┌──────────────┐ ┌───────────────┐ ┌────────────┐ │\n│ │ CuTeDSL kernel │ │ cute.compile()│ │.cu │ │.obj │ │\n│ │ (@cute.jit + │──▶│ export_to_c()│──▶│(struct init + │──▶│(from nvcc) │ │\n│ │ @cute.kernel) │ │ │ │ metadata load)│ │ │ │\n│ └─────────────────┘ └──────┬───────┘ └───────────────┘ └────┬───────┘ │\n│ │ │ link │\n│ ┌─────▼────────┐ ┌─────▼───────┐ │\n│ │_cutedsl.h │ │ test.so │ │\n│ │_cutedsl.o │────────────────────▶│ (+libcuda) │ │\n│ │(pre-compiled) │ └─────────────┘ │\n│ └──────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n```\n\n### Key Difference\n\n- **CUTLASS C++ backend**: Jinja2 template IS the kernel. It generates C++ source\n that instantiates CUTLASS template classes. `nvcc` compiles everything.\n- **CuTeDSL backend**: Kernel is written in Python using CuTeDSL decorators\n (`@cute.jit`, `@cute.kernel`). `cute.compile()` produces MLIR+cubin at\n codegen time, `export_to_c()` emits `.h`+`.o`. Only a thin C++ wrapper\n needs nvcc compilation.\n\n---\n\n## Step-by-Step Integration Guide\n\n### Step 1: Write the CuTeDSL Kernel\n\n**Create**: `backend/cuda//cutedsl__sm80.py`\n\nThis file contains the pure CuTeDSL kernel implementation:\n\n```python\nimport cutlass\nimport cutlass.cute as cute\nfrom cutlass.cute.nvgpu import cpasync, warp\n\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n self.param = param\n self.cta_sync_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=128)\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, mB: cute.Tensor, ..., stream: cuda.CUstream):\n \"\"\"Host-side: create SMEM layouts, copy atoms, MMA config, launch kernel.\"\"\"\n # 1. Compute swizzled SMEM layouts\n sA_layout = cute.tile_to_shape(get_smem_layout_atom(...), ...)\n # 2. Create copy atoms (cpasync for GMEM->SMEM, universal for stores)\n gmem_copy = cute.make_tiled_copy_tv(cute.make_copy_atom(cpasync.CopyG2SOp(), ...), ...)\n # 3. Create tiled MMA\n tiled_mma = cute.make_tiled_mma(warp.MmaF16BF16Op(...), ...)\n # 4. Compute grid dimensions\n grid_dim = (m_blocks, batch_size, 1)\n # 5. Launch kernel\n self.kernel(...).launch(grid=grid_dim, block=[128, 1, 1], stream=stream)\n\n @cute.kernel\n def kernel(self, ...all_params...):\n \"\"\"Device-side: actual CUDA kernel body.\"\"\"\n tidx, _, _ = cute.arch.thread_idx()\n m_block, batch, _ = cute.arch.block_idx()\n smem = cutlass.utils.SmemAllocator()\n # ... GEMM mainloop, epilogue, store ...\n```\n\n**Key patterns**:\n\n1. **Swizzled SMEM layouts** (for bank-conflict-free access):\n ```python\n def get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = (128 if bytes_per_row % 128 == 0 else ...) // dtype_byte\n swizzle_bits = 4 if smem_k_block_size == 128 else ...\n return cute.make_composed_layout(\n cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), 0,\n cute.make_ordered_layout((8, smem_k_block_size), order=(1, 0)),\n )\n ```\n\n2. **SM80 warp-level MMA** (`warp.MmaF16BF16Op`):\n ```python\n tiled_mma = cute.make_tiled_mma(\n warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)),\n (num_warps, 1, 1),\n permutation_mnk=(num_warps * 16, 16, 16),\n )\n ```\n\n3. **cp.async for GMEM->SMEM**:\n ```python\n atom_async = cute.make_copy_atom(\n cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),\n dtype, num_bits_per_copy=128)\n gmem_tiled_copy = cute.make_tiled_copy_tv(atom_async, thread_layout, value_layout)\n ```\n\n4. **SMEM->register via ldmatrix**:\n ```python\n smem_copy_atom = cute.make_copy_atom(\n warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), dtype)\n smem_tiled_copy_A = cute.make_tiled_copy_A(smem_copy_atom, tiled_mma)\n ```\n\n5. **Activation dispatch table**:\n ```python\n ACTIVATION_MAP = {\n \"Sigmoid\": lambda x: 1.0 / (1.0 + cute.math.exp(-x)),\n \"ReLu\": lambda x: cute.arch.fmax(x, 0.0),\n \"Identity\": lambda x: x,\n }\n ```\n\n### Step 2: Create the Backend Codegen Module\n\n**Create**: `backend/cuda//_cutedsl.py`\n\nThis file registers AIT backend functions with `_cutedsl` suffix:\n\n```python\nimport functools, os, jinja2\nfrom aitemplate.backend import registry\nfrom aitemplate.backend.target import Target\n\n# C++ wrapper template - constructs CuTeDSL typed tensor structs from AIT void* pointers\nCUTEDSL_WRAPPER_TEMPLATE = jinja2.Template(\"\"\"\n#include \n#include \"{{cutedsl_header}}\"\n\nstatic {{func_name}}_cutedsl_Kernel_Metadata_t g_metadata;\nstatic bool g_metadata_loaded = false;\n\nstatic void ensure_metadata_loaded() {\n if (!g_metadata_loaded) {\n cuInit(0);\n {{func_name}}_cutedsl_Kernel_Metadata_Load(&g_metadata);\n g_metadata_loaded = true;\n }\n}\n\nvoid {{func_name}}(...AIT signature...) {\n ensure_metadata_loaded();\n\n // Construct CuTeDSL typed tensor structs:\n {{func_name}}_cutedsl_Tensor_mA_t t_mA;\n t_mA.data = a_ptr;\n t_mA.dynamic_shapes[0] = batch; // first dynamic dim\n t_mA.dynamic_shapes[1] = seq_len; // second dynamic dim\n t_mA.dynamic_strides[0] = seq_len * K; // batch stride\n\n // ... same for all tensors ...\n\n CUstream cu_stream = reinterpret_cast(stream);\n {{func_name}}_cutedsl_wrapper(&g_metadata, &t_mA, ..., cu_stream);\n}\n\"\"\")\n\n@functools.lru_cache(maxsize=32)\ndef _aot_compile_cutedsl_kernel(n0, n1, ..., output_dir, func_name, arch):\n \"\"\"AOT compile kernel and export .h + .o.\"\"\"\n from my_kernel import MySm80Kernel\n from cutlass.cute.runtime import from_dlpack\n from cutlass.cute.export import export_to_c\n\n kernel = MySm80Kernel(n0=n0, n1=n1, ...)\n\n # Create representative tensors (shapes don't matter for dynamic dims)\n q_pt = torch.zeros(4, 128, n1, device=\"cuda\", dtype=torch.float16)\n\n # Mark dynamic modes with per-mode divisibility\n q_cute = from_dlpack(q_pt, assumed_align=16)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=0, stride_order=q_pt.dim_order(), divisibility=1)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=1, stride_order=q_pt.dim_order(), divisibility=8)\n\n # Compile and export\n compiled = cute.compile(kernel, q_cute, ..., cu_stream)\n export_to_c(compiled, file_path=output_dir, file_name=f\"{func_name}_cutedsl\")\n return h_path, o_path\n\n@registry.reg(\"cuda..gen_function_cutedsl\")\ndef gen_function_cutedsl(func_attrs):\n h_path, o_path = _aot_compile_cutedsl_kernel(...)\n func_attrs[\"cutedsl_obj_path\"] = o_path # For builder to pick up\n return CUTEDSL_WRAPPER_TEMPLATE.render(...)\n\n@registry.reg(\"cuda..func_decl_cutedsl\")\ndef func_decl_cutedsl(func_attrs): ...\n\n@registry.reg(\"cuda..func_call_cutedsl\")\ndef func_call_cutedsl(func_attrs, indent=\" \"): ...\n```\n\n**Critical detail**: The C++ wrapper MUST construct CuTeDSL typed tensor structs\n(`*_Tensor_mX_t`) with proper `dynamic_shapes[]` and `dynamic_strides[]` fields.\nThe `export_to_c()` API generates these structs -- they are NOT raw `void*`.\n\n### Step 3: Add Backend Dispatch in the Op Class\n\n**Modify**: `compiler/ops//.py`\n\nAdd methods to detect CuTeDSL mode and dispatch with suffix:\n\n```python\ndef _use_cutedsl(self) -> bool:\n current_target = target.Target.current()\n return current_target._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n\ndef gen_function(self) -> str:\n suffix = self._backend_suffix()\n self._attrs[\"backend_suffix\"] = suffix # For codegen.py to pick up\n func_key = f\"{target_name}.{op_name}.gen_function{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_decl(self, func_attrs=None) -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_decl{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_call(self, func_attrs=None, indent=\" \") -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_call{suffix}\"\n return registry.get(func_key)(self._attrs, indent)\n```\n\n### Step 4: Modify the Build System\n\n#### 4.1 `backend/codegen.py` -- `gen_function_src()`\n\nPass `workdir` to `func_attrs` so the AOT compiler knows where to write:\n```python\nfunc._attrs[\"workdir\"] = prefix\n```\n\nAfter writing the `.cu` wrapper, add the pre-compiled `.o` to `file_pairs`:\n```python\ncutedsl_obj = func._attrs.get(\"cutedsl_obj_path\")\nif cutedsl_obj and os.path.exists(cutedsl_obj):\n file_pairs.append((cutedsl_obj, cutedsl_obj))\n```\n\nIn `_process_src_ops()`, read `backend_suffix` for `func_decl` / `func_call` dispatch:\n```python\nbackend_suffix = func._attrs.get(\"backend_suffix\", \"\")\nf_func_decl = registry.get(f\"{target}.{op}.func_decl{backend_suffix}\")\nf_func_call = registry.get(f\"{target}.{op}.func_call{backend_suffix}\")\n```\n\n#### 4.2 `backend/builder.py` -- Handle pre-compiled `.o` files\n\nIn `build_objs()`: skip compilation for `.o` source files:\n```python\nif src.endswith(\".o\"):\n _LOGGER.info(f\"Skipping compilation for pre-compiled object: {src}\")\n continue\n```\n\nIn `build_so()`: add `-lcuda` when CuTeDSL objects are present:\n```python\nhas_cutedsl = any(obj.endswith(\"_cutedsl.o\") for obj in objs)\nextra_libs = \" -lcuda\" if has_cutedsl else \"\"\n```\n\nIn `gen_makefile()`: use `os.path.basename()` for `.o` paths (Make runs from workdir):\n```python\nif pair[0].endswith(\".o\"):\n obj_files.append(os.path.basename(pair[0])) # NOT full path\n has_cutedsl = True\n```\n\n### Step 5: Update Target Configuration\n\n**Modify**: `backend/target.py`\n\nAdd `use_cutedsl_` kwarg support to the Target class (already passed through\n`_kwargs`; no explicit changes needed unless you want validation).\n\n### Step 6: Update `__init__.py` and Examples\n\n**Modify**: `backend/cuda//__init__.py`\n\nAdd the new module import for auto-registration:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n**Modify**: Example/test scripts\n\nAdd CLI arg to enable CuTeDSL:\n```python\ntarget = Target(..., use_cutedsl_=True)\n```\n\n### Step 7: Add BUCK Dependencies\n\nAdd `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps\nin `examples/BUCK` for any target that uses the CuTeDSL backend.\n\n---\n\n## Common Pitfalls and Debug Checklist\n\n| # | Error | Root Cause | Fix |\n|---|-------|------------|-----|\n| 1 | `ModuleNotFoundError: No module named 'cuda'` | Missing `nvidia-cutlass-dsl` BUCK dep | Add `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps |\n| 2 | `ValueError: Invalid leading dimension: 2` | `LayoutEnum.from_tensor()` only accepts 2D tensors (leading_dim 0 or 1) | Hardcode `utils.LayoutEnum.ROW_MAJOR` for 3D row-major tensors |\n| 3 | `_ScaledBasis TypeError` in SM90 TMA | CuTeDSL TMA layout algebra doesn't yet handle 3D tensors well | Use SM80 kernel (works on SM90+ hardware too) |\n| 4 | `coord and shape of view are weakly congruent: N modes vs M coords` | After `partition_S(local_tile(..., (tile, None)))`, the tile mode creates an extra dimension | Add extra `None` to indexing: `tQgQ[None, None, None, k_tile]` instead of `tQgQ[None, None, k_tile]` |\n| 5 | `ptr alignment does not meet requirement` for 128-bit cp.async | `mark_layout_dynamic()` loses alignment info | Use `mark_compact_shape_dynamic(mode=M, divisibility=8)` to preserve alignment |\n| 6 | `cute.copy expects src and dst to have static shape` | `mark_layout_dynamic()` makes ALL dims dynamic, but inner dims must be static constants | Use per-mode `mark_compact_shape_dynamic()` only on truly dynamic modes (batch, seq_len); leave inner dims (N0, N1, K) static |\n| 7 | `shape(X) of mode(0) is not divisible by divisibility(8)` | Batch dim size not divisible by the divisibility requirement | Use per-mode divisibility: batch `div=1`, M/seq_len `div=8` |\n| 8 | `argument of type \"void *\" incompatible with \"..._Tensor_mX_t *\"` | CuTeDSL `export_to_c` generates typed tensor structs, not raw pointer APIs | Construct typed tensor structs in C++ wrapper: set `.data`, `.dynamic_shapes[]`, `.dynamic_strides[]` |\n| 9 | `No rule to make target 'tmp/.../_cutedsl.o'` | Makefile has full path but Make runs from workdir | Use `os.path.basename(pair[0])` in `gen_makefile()` |\n| 10 | `DT_TEXTREL` linker warning | CuTeDSL `.o` contains text relocations | Non-fatal warning; can be ignored |\n\n---\n\n## Dynamic Tensor Marking Strategy\n\nThis is the most critical part for correctness. CuTeDSL needs to know which\ntensor dimensions are compile-time constants vs runtime-dynamic.\n\n### DO NOT use `mark_layout_dynamic()`\n\n`mark_layout_dynamic(leading_dim=N)` makes ALL shapes and strides dynamic.\nThis breaks:\n- `cute.copy` which requires static shapes for src/dst\n- cp.async alignment checks (loses divisibility info)\n- Optimal code generation (no compile-time constants)\n\n### DO use `mark_compact_shape_dynamic()` per mode\n\n```python\ndef make_cute_tensor(torch_tensor, dynamic_modes_div):\n \"\"\"\n dynamic_modes_div: list of (mode_index, divisibility) tuples.\n Only the specified modes become dynamic; all other dims stay static.\n \"\"\"\n ct = from_dlpack(torch_tensor, assumed_align=16)\n for mode, div in dynamic_modes_div:\n ct = ct.mark_compact_shape_dynamic(\n mode=mode,\n stride_order=torch_tensor.dim_order(),\n divisibility=div,\n )\n return ct\n```\n\n**Typical b2b_bmm pattern**:\n| Tensor | Shape | Dynamic Modes |\n|--------|-------|---------------|\n| Q | [batch, M, K] | batch (div=1), M (div=8) |\n| K | [batch, N0, K] | batch (div=1) only |\n| V | [batch, N0, N1] | batch (div=1) only |\n| Bias | [batch, M, N0] | batch (div=1), M (div=8) |\n| Out | [batch, M, N1] | batch (div=1), M (div=8) |\n\n### Why div=8 for M?\n\n128-bit cp.async copies 8 fp16 elements at once. The M dimension appears in\nstrides that feed cp.async, so CuTeDSL needs `divisibility=8` to guarantee\nalignment for 128-bit vectorized copies.\n\n---\n\n## Generated Artifact Reference\n\nAfter successful compilation, the workdir contains:\n\n| File | Size (typical) | Description |\n|------|----------------|-------------|\n| `_cutedsl.h` | ~4 KB | CuTeDSL header: typed tensor structs, `Metadata_Load/Unload`, `wrapper()` |\n| `_cutedsl.o` | ~50 KB | Pre-compiled object with embedded cubin(s) |\n| `.cu` | ~4 KB | Thin C++ wrapper: `cuInit`, metadata load, struct construction, launch |\n| `.obj` | ~16 KB | Compiled wrapper object |\n| `test.so` | ~1.4 MB | Final shared library (linked with `-lcuda`) |\n\n### Header structure (`_cutedsl.h`)\n\n```c\n// Metadata struct (holds CUlibrary handle)\ntypedef struct {\n CUlibrary module;\n} _cutedsl_Kernel_Metadata_t;\n\n// Typed tensor structs (one per kernel argument)\ntypedef struct {\n void *data;\n int32_t dynamic_shapes[N_DYNAMIC_SHAPES]; // runtime batch, seq_len, etc.\n int64_t dynamic_strides[N_DYNAMIC_STRIDES]; // runtime strides\n} _cutedsl_Tensor_mX_t;\n\n// Load/Unload/Wrapper functions\nvoid _cutedsl_Kernel_Metadata_Load(_cutedsl_Kernel_Metadata_t *metadata);\nvoid _cutedsl_Kernel_Metadata_Unload(_cutedsl_Kernel_Metadata_t *metadata);\nint32_t _cutedsl_wrapper(_cutedsl_Kernel_Metadata_t *metadata,\n _cutedsl_Tensor_mX_t *mX, ..., CUstream stream);\n```\n\n---\n\n## File Checklist (for classic_b2b_bmm)\n\n### Files Created\n\n| File | Path (relative to aitemplate/) |\n|------|-------------------------------|\n| SM80 kernel | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n| SM90 kernel (WIP) | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm90.py` |\n| Backend codegen | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| AOT example | `examples/07_how_to_run_pt_model/cutedsl_b2b_bmm_aot_example.py` |\n\n### Files Modified\n\n| File | Change |\n|------|--------|\n| `compiler/ops/b2b_bmm/classic_b2b_bmm.py` | Added `_use_cutedsl()`, `_backend_suffix()`, dispatch in `gen_function/decl/call` |\n| `backend/codegen.py` | Pass `workdir`, handle `.o` in `file_pairs`, read `backend_suffix` |\n| `backend/builder.py` | Skip `.o` compilation, add `-lcuda`, use `basename()` in Makefile |\n| `backend/cuda/b2b_bmm/__init__.py` | Added `classic_b2b_bmm_cutedsl` import |\n| `examples/07_how_to_run_pt_model/classic_b2b_bmm_example.py` | Added `--use-cutedsl` and `--both` CLI args |\n| `examples/BUCK` | Added `nvidia-cutlass-dsl` dependency |\n\n---\n\n## Reference Codebases\n\n| Reference | Path | What to reuse |\n|-----------|------|---------------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | SM80 warp MMA, SMEM layout, gemm/gemm_rs patterns |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` | SM90 TMA+WGMMA, pipeline patterns |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` | Complete CuTeDSL GEMM lifecycle |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` | `export_to_c()`, `dump_to_object()` API |\n| Quack PT2 compat | `fbcode/ads_mkl/ops/cute_dsl/quack/quack_rmsnorm_pt2.py` | `@functools.cache` + `cute.compile()` caching |\n| Inductor CuTeDSL | `fbcode/caffe2/torch/_inductor/codegen/cutedsl/` | Framework integration pattern |\n| Current CUTLASS backend | `backend/cuda/b2b_bmm/classic_b2b_bmm.py` | Interface contract (signature, registry keys) |\n| CUTLASS C++ kernels | `static/include/kernels/classic_b2b_bmm/` | Algorithm reference for register-passing b2b GEMM |\n\n---\n\n## Claude Skill: `add-cutedsl-backend`\n\nBelow is a reusable Claude skill specification for adding a CuTeDSL backend\nto any AIT operator.\n\n### Skill Description\n\n```\nWhen adding a CuTeDSL backend for an AITemplate operator, follow these steps:\n\n1. RESEARCH: Read the existing CUTLASS C++ backend for the target op:\n - `backend/cuda//.py` -- understand FUNC_TEMPLATE, registry keys,\n tensor shapes, tiling parameters\n - `compiler/ops//.py` -- understand gen_function/decl/call dispatch\n\n2. KERNEL: Write the CuTeDSL kernel in `backend/cuda//cutedsl__sm80.py`:\n - Class with __init__ (compile-time params), @cute.jit __call__ (host setup),\n @cute.kernel kernel (device code)\n - SM80: use warp.MmaF16BF16Op, cpasync.CopyG2SOp, warp.LdMatrix8x8x16bOp\n - Swizzled SMEM layouts via get_smem_layout_atom()\n - Named barriers for CTA sync\n\n3. CODEGEN: Create `backend/cuda//_cutedsl.py`:\n - @functools.lru_cache AOT compile function\n - C++ wrapper template with typed tensor struct construction\n - Registry functions with _cutedsl suffix\n - CRITICAL: Use mark_compact_shape_dynamic() per-mode, NOT mark_layout_dynamic()\n - CRITICAL: C++ wrapper must construct *_Tensor_mX_t structs, NOT pass void*\n\n4. DISPATCH: Modify `compiler/ops//.py`:\n - Add _use_cutedsl(), _backend_suffix()\n - Add suffix dispatch in gen_function(), gen_function_decl(), gen_function_call()\n - Store backend_suffix in self._attrs\n\n5. BUILD: Modify `backend/codegen.py` and `backend/builder.py`:\n - codegen.py: pass workdir, handle .o in file_pairs, read backend_suffix\n - builder.py: skip .o compilation, add -lcuda, use basename() for .o in Makefile\n\n6. REGISTER: Add import in `backend/cuda//__init__.py`\n\n7. TEST: Run with Target(..., use_cutedsl_=True) and validate against\n PyTorch reference (atol=1e-2 for fp16)\n\nCommon pitfalls to watch for:\n- mark_layout_dynamic makes ALL dims dynamic -> breaks static copy requirements\n- partition_S(local_tile(..., (tile, None))) adds extra modes -> count None indices\n- CuTeDSL export_to_c generates typed structs, not raw void* -> read .h header\n- Makefile .o paths must use basename (Make runs from workdir)\n- Need -lcuda for CUlibrary-based metadata loading\n```\n\n---\n\n## Execution Order for New Operators\n\n1. **Read existing backend** -- understand tensor shapes, tiling, registry keys\n2. **Write SM80 kernel** -- implement, test standalone with `cute.compile()`\n3. **AOT export test** -- verify `export_to_c()` produces valid `.h` + `.o`\n4. **Create backend codegen** -- registry functions, C++ wrapper template\n5. **Add dispatch** -- op class modifications, `__init__.py` import\n6. **Modify build system** -- codegen.py, builder.py changes\n7. **End-to-end test** -- run full AIT pipeline, compare outputs\n8. **(Optional) SM90 kernel** -- TMA + WGMMA version for Hopper\n\n---\n\n## Verification Checklist\n\n- [ ] CuTeDSL kernel compiles: `cute.compile()` succeeds\n- [ ] AOT export produces `.h` and `.o` files\n- [ ] C++ wrapper compiles with nvcc (no type mismatches)\n- [ ] `.o` links into `.so` (no undefined symbols, `-lcuda` present)\n- [ ] Runtime loads metadata: `Kernel_Metadata_Load()` succeeds\n- [ ] Numerical correctness: output matches PyTorch reference (atol<=0.02 for fp16)\n- [ ] Dynamic shapes work: different batch sizes / sequence lengths\n- [ ] Original CUTLASS backend still works (no regression)\n","isInternal":false,"tokens":6479,"sizeBytes":24922},{"name":"migrate-gemm-to-cutedsl.md","path":".claude/skills/migrate-gemm-to-cutedsl.md","rawUrl":"https://raw.githubusercontent.com/facebookincubator/AITemplate/HEAD/.claude/skills/migrate-gemm-to-cutedsl.md","title":"Skills Skill","category":"anthropic-skill","format":"markdown","content":"## Migrating GEMM Operators from CUTLASS C++ to CuTeDSL\n\n### Overview\n\nThis skill documents the pattern for migrating any GEMM variant (gemm_rcr_bias, gemm_rrr, gemm_ccr, etc.) from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS Python DSL), using AOT-style deployment.\n\n### Layout Reference\n\nCUTLASS GEMM naming convention: `gemm_` where R=Row-major, C=Column-major.\n\n| Layout | A shape/order | B shape/order | MMA Atom | Key property |\n|--------|--------------|--------------|----------|-------------|\n| RCR | [M,K] K-major | [N,K] K-major | TN | Both K-contiguous |\n| RRR | [M,K] K-major | [K,N] N-major | TN/NT | A=K-cont, B=N-cont |\n| CCR | [K,M] M-major | [N,K] K-major | NN | A=M-cont, B=K-cont |\n\nFor SM80 with fp16: `warp.MmaF16BF16Op(Float16, Float32, (16,8,16))` with TN layout.\n\n### Step-by-Step Migration\n\n#### 1. Identify the operation\n\nRead the CUTLASS backend file (e.g., `backend/cuda/gemm_universal/gemm_rcr_bias.py`) to understand:\n- The GEMM layout (RCR, RRR, etc.)\n- The epilogue (bias, relu, bias+relu, etc.)\n- The function signature from `common_bias.py` or `common.py`\n- Whether it supports split-K\n\n#### 2. Create the CuTeDSL kernel\n\nCreate `backend/cuda/gemm_universal/cutedsl__sm80.py`.\n\nKey components to implement:\n\n```python\nclass GemmXxxSm80Kernel:\n def __init__(self, tile_m=128, tile_n=128, tile_k=32):\n self.tile_m, self.tile_n, self.tile_k = tile_m, tile_n, tile_k\n self.num_threads = 128 # 4 warps\n\n @cute.jit\n def __call__(self, mA, mB, mC, ...):\n # 1. Set up SMEM layouts using get_smem_layout_atom()\n # 2. Define SharedStorage struct\n # 3. Create copy atoms (async for load, universal for store)\n # 4. Create tiled_mma with SM80 16x8x16 atom\n # 5. Compute grid dims and launch kernel\n\n @cute.kernel\n def kernel(self, ...):\n # 1. Allocate SMEM, partition GMEM tiles\n # 2. GEMM mainloop: K-loop with cp.async + warp MMA\n # 3. Epilogue: apply bias/activation on FP32 accumulators\n # 4. Store: Reg -> SMEM -> GMEM\n```\n\n**SMEM layout pattern** (from `cutedsl_b2b_bmm_sm80.py`):\n```python\ndef get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = cutlass.const_expr(\n 128 if bytes_per_row % 128 == 0\n else (64 if bytes_per_row % 64 == 0\n else (32 if bytes_per_row % 32 == 0 else 16))\n ) // dtype_byte\n # ... swizzle computation\n return cute.make_composed_layout(swizzle, 0, ordered_layout)\n```\n\n**GEMM mainloop pattern**:\n```python\nfor k_tile in range(k_tiles):\n cute.copy(gmem_tiled_copy, src_gmem, dst_smem)\n cute.arch.cp_async_commit_group()\n cute.arch.cp_async_wait_group(0)\n barrier.arrive_and_wait()\n\n # Load first k-block from SMEM to registers\n cute.copy(smem_tiled_copy_A, tSsA[None, None, 0], tSrA_view[None, None, 0])\n cute.copy(smem_tiled_copy_B, tSsB[None, None, 0], tSrB_view[None, None, 0])\n\n for k_block in cutlass.range_constexpr(num_k_blocks):\n k_next = (k_block + 1) % num_k_blocks\n cute.copy(smem_tiled_copy_A, tSsA[None, None, k_next], ...)\n cute.copy(smem_tiled_copy_B, tSsB[None, None, k_next], ...)\n cute.gemm(tiled_mma, acc, tSrA[None, None, k_block], tSrB[None, None, k_block], acc)\n\n barrier.arrive_and_wait()\n```\n\n**Store pattern** (Reg -> SMEM -> GMEM):\n```python\nrOut = cute.make_fragment_like(acc, dtype)\nrOut.store(acc.load().to(dtype))\nsOut = cute.make_tensor(sA.iterator, sOut_layout) # reuse SMEM\n\nsmem_copy_atom_C = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype)\nsmem_tiled_copy_C = cute.make_tiled_copy_C(smem_copy_atom_C, tiled_mma)\n# ... copy from reg to SMEM, sync, then SMEM to GMEM\n```\n\n#### 3. Handle the epilogue\n\n**1D bias (broadcast)**: Load bias into SMEM as a 2D view with stride-0 in M, use `thr_mma.partition_C()` to align with accumulator layout:\n```python\n# In __call__: create 2D broadcast view\nmBias_2d = cute.make_tensor(mBias.iterator,\n cute.make_layout((M_size, N_size), stride=(0, bias_stride)))\n\n# In kernel: tile and load\ngBias = cute.local_tile(mBias_2d, (tile_m, tile_n), (m_block, n_block))\n# ... cp.async load to sBias ...\ntCsBias = thr_mma.partition_C(sBias)\n# element-wise add in FP32\n```\n\n**2D bias**: Direct tile and load (like b2b_bmm does).\n\n**Activation**: Apply element-wise in FP32 accumulator space before store.\n\n#### 4. Create AOT example\n\nCreate `examples/07_how_to_run_pt_model/cutedsl__aot_example.py`:\n\n```python\ndef make_cute_tensor(torch_tensor):\n ct = from_dlpack(torch_tensor, assumed_align=16)\n ct = ct.mark_layout_dynamic(leading_dim=len(torch_tensor.shape) - 1)\n ct = ct.mark_compact_shape_dynamic(\n mode=len(torch_tensor.shape) - 1,\n stride_order=torch_tensor.dim_order(),\n divisibility=(128 // cutlass.Float16.width),\n )\n return ct\n\ndef jit_compile_and_validate():\n # 1. Create PyTorch reference tensors\n # 2. Compute PyTorch reference output\n # 3. Create kernel instance\n # 4. cute.compile(kernel, *cute_tensors, *scalar_args, stream)\n # 5. Execute and compare (atol=1e-2, rtol=1e-2 for fp16)\n\ndef aot_export(compiled):\n # export_to_c(compiled, file_path=output_dir, file_name=name)\n # Generates: .h + .o\n```\n\n#### 5. Add BUCK target\n\n```python\npython_binary(\n # @autodeps-skip\n name = \"cutedsl__aot_example\",\n srcs = [\"07_how_to_run_pt_model/cutedsl__aot_example.py\"],\n main_module = \"aitemplate.AITemplate.examples.07_how_to_run_pt_model.cutedsl__aot_example\",\n deps = [\n \"fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl\",\n \"//aitemplate/AITemplate/python/aitemplate:aitemplate\",\n \"//caffe2:torch\",\n ],\n)\n```\n\n### Reference Implementations\n\n| File | What it demonstrates |\n|------|---------------------|\n| `backend/cuda/gemm_universal/cutedsl_gemm_rcr_bias_sm80.py` | GEMM RCR + 1D bias, 2D grid |\n| `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` | Fused B2B GEMM, 2D bias, activation |\n| `examples/.../cutedsl_gemm_rcr_bias_aot_example.py` | AOT example for GEMM |\n| `examples/.../cutedsl_b2b_bmm_aot_example.py` | AOT example for B2B BMM |\n| `ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | Reusable SM80 helpers |\n\n### Common Pitfalls\n\n1. **Alignment**: Use `mark_compact_shape_dynamic(divisibility=8)` for fp16 to satisfy 128-bit cp.async alignment\n2. **Bias broadcast**: Create 2D view with stride-0 in broadcast dim at host level (`__call__`), not inside `@cute.kernel`\n3. **SMEM reuse**: Output staging can reuse SMEM from input tiles (e.g., `sC = cute.make_tensor(sA.iterator, sC_layout)`)\n4. **MN view**: Use `_make_acc_tensor_mn_view()` to iterate accumulator in logical (M, N) order for epilogue ops\n5. **Python slicing**: CuTe tensors don't support Python slicing (`tensor[start:]`). Use `cute.local_tile()` or `cute.make_tensor()` with offset iterators\n","isInternal":false,"tokens":2134,"sizeBytes":7026},{"name":"testing.md","path":".claude/skills/testing.md","rawUrl":"https://raw.githubusercontent.com/facebookincubator/AITemplate/HEAD/.claude/skills/testing.md","title":"Skills Skill","category":"anthropic-skill","format":"markdown","content":"## Testing in AITemplate\n\n### Running Tests\n\n#### Python tests with pytest\n```bash\n# Run a specific test file\npython -m pytest tests/unittest/test_my_op.py -v\n\n# Run a specific test case\npython -m pytest tests/unittest/test_my_op.py::TestMyOp::test_basic -v\n```\n\n#### Buck tests (Meta-internal)\n```bash\n# Run a specific Buck test target\nbuck2 test //aitemplate/AITemplate:test_target_name\n```\n\n### Test Patterns\n\nAITemplate tests typically follow this structure:\n\n```python\nimport unittest\nimport torch\nfrom aitemplate.testing import detect_target\nfrom aitemplate.compiler import compile_model\nfrom aitemplate.frontend import Tensor, nn\n\nclass TestMyFeature(unittest.TestCase):\n def test_basic(self):\n # 1. Check for available GPU target\n target = detect_target()\n\n # 2. Build AIT graph\n X = Tensor(shape=[2, 3], name=\"X\", is_input=True)\n Y = my_op()(X)\n Y._attrs[\"is_output\"] = True\n Y._attrs[\"name\"] = \"Y\"\n\n # 3. Compile\n module = compile_model(Y, target, \"./tmp\", \"test_my_op\")\n\n # 4. Run with real data\n x_pt = torch.randn(2, 3).cuda().half()\n y = torch.empty(2, 3).cuda().half()\n module.run_with_tensors([x_pt], [y])\n\n # 5. Compare against PyTorch reference\n y_ref = torch_reference_impl(x_pt)\n torch.testing.assert_close(y, y_ref)\n```\n\n### Key Testing Utilities\n\n- `aitemplate.testing.detect_target()` — auto-detect CUDA or ROCm\n- `compile_model()` — compile a graph to a loadable module\n- `module.run_with_tensors()` — execute the compiled model\n- Tests typically use `float16` precision and compare with tolerance\n\n### FX2AIT Tests\n\nFor the PyTorch FX converter, tests are in `fx2ait/fx2ait/test/`:\n```python\n# These tests verify that PyTorch models convert correctly to AIT\n# and produce numerically equivalent results\n```\n","isInternal":false,"tokens":479,"sizeBytes":1853}],"systemPromptSnippet":"\n\n# Skill: Add CuTeDSL Backend to an AITemplate Operator\n\n## Description\n\nThis skill guides adding a CuTeDSL (CUTLASS Python DSL) backend to an existing\nAITemplate operator, replacing the Jinja2/CUTLASS C++ code generation with\nPython-based AOT compilation via `cute.compile()` + `export_to_c()`.\n\n## When to Use\n\n- Converting an existing AIT operator from CUTLASS C++ templates to CuTeDSL\n- Adding an alternative CuTeDSL backend alongside the existing CUTLASS backend\n- Creating a new GPU kernel using CuTeDSL within the AIT framework\n\n## Prerequisites\n\n- The operator already has a working CUTLASS C++ backend\n- `nvidia-cutlass-dsl` package is available in the build system\n- Target GPU is SM80+ (Ampere or newer)\n\n## Steps\n\n### 1. RESEARCH the existing backend\n\nRead these files to understand the operator interface:\n\n```\nbackend/cuda//.py # FUNC_TEMPLATE, registry keys, tensor shapes\ncompiler/ops//.py # gen_function/decl/call dispatch\n```\n\nKey things to note:\n- Function signature (what void* params, what scalar dims)\n- Tensor shapes and which dims are static vs dynamic\n- Tiling parameters (threadblock M/N/K, warp M/N/K)\n- Registry key pattern (e.g., `cuda.classic_b2b_bmm.gen_function`)\n\n### 2. WRITE the CuTeDSL kernel\n\nCreate: `backend/cuda//cutedsl__sm80.py`\n\nStructure:\n```python\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n # Store tile sizes, compile-time constants\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, ..., stream: cuda.CUstream):\n # Host-side: SMEM layouts, copy atoms, MMA config, kernel launch\n\n @cute.kernel\n def kernel(self, ...):\n # Device-side: actual CUDA kernel body\n```\n\nKey SM80 patterns:\n- `warp.MmaF16BF16Op` for MMA atoms\n- `cpasync.CopyG2SOp()` for GMEM->SMEM async copy\n- `warp.LdMatrix8x8x16bOp` for SMEM->register loads\n- Swizzled SMEM layouts via `cute.make_composed_layout(cute.make_swizzle(...))`\n- `pipeline.NamedBarrier` for CTA synchronization\n\n### 3. CREATE the backend codegen module\n\nCreate: `backend/cuda//_cutedsl.py`\n\nMust contain:\n- `@functools.lru_cache` AOT compilation function\n- C++ wrapper template (Jinja2) that constructs typed tensor structs\n- Three registry functions with `_cutedsl` suffix:\n - `cuda..gen_function_cutedsl`\n - `cuda..func_decl_cutedsl`\n - `cuda..func_call_cutedsl`\n\nCRITICAL: The C++ wrapper must construct `*_Tensor_mX_t` typed structs\n(not pass raw void*). Read the generated `.h` header to see struct fields.\n\nCRITICAL: Use `mark_compact_shape_dynamic()` per-mode, NOT `mark_layout_dynamic()`.\nOnly mark truly dynamic dims (batch, seq_len). Keep inner dims static.\n\n### 4. ADD backend dispatch\n\nModify: `compiler/ops//.py`\n\n```python\ndef _use_cutedsl(self) -> bool:\n return target.Target.current()._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n```\n\nAdd suffix dispatch in `gen_function()`, `gen_function_decl()`, `gen_function_call()`.\nStore `backend_suffix` in `self._attrs` for codegen.py to pick up.\n\n### 5. MODIFY the build system\n\n`backend/codegen.py`:\n- Pass `workdir` to `func_attrs`\n- Add CuTeDSL `.o` to `file_pairs`\n- Read `backend_suffix` for func_decl/func_call dispatch\n\n`backend/builder.py`:\n- Skip compilation for pre-compiled `.o` files\n- Add `-lcuda` when CuTeDSL objects are present\n- Use `os.path.basename()` for `.o` paths in Makefile\n\n### 6. REGISTER the new module\n\nAdd import in `backend/cuda//__init__.py`:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n### 7. TEST\n\nRun with `Target(..., use_cutedsl_=True)`.\nValidate against PyTorch reference (atol=1e-2 for fp16).\n\n## Common Pitfalls\n\n1. `mark_layout_dynamic()` makes ALL dims dynamic -> breaks static copy reqs\n2. `partition_S(local_tile(..., (tile, None)))` adds extra modes -> count None indices\n3. `export_to_c` generates typed structs, not raw void* -> read .h header\n4. Makefile .o paths must use basename (Make runs from workdir)\n5. Need `-lcuda` for CUlibrary-based metadata loading\n6. Batch dim divisibility: use div=1 for batch, div=8 for seq_len\n7. `LayoutEnum.from_tensor()` only handles 2D -> hardcode ROW_MAJOR for 3D\n\n## References\n\n| Reference | Path |\n|-----------|------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` |\n| classic_b2b_bmm impl | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| classic_b2b_bmm SM80 | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n\n\n## Adding New Operators to AITemplate\n\n### Overview\n\nAdding a new operator involves defining it in Python, writing backend codegen, and adding tests.\n\n### Steps\n\n#### 1. Define the Operator (Python)\n\nCreate a new file in `python/aitemplate/compiler/ops/` (or the appropriate subdirectory):\n\n```python\n# python/aitemplate/compiler/ops/common/my_new_op.py\n\nfrom aitemplate.compiler.base import Operator, Tensor\n\nclass my_new_op(Operator):\n \"\"\"Description of what the op does.\"\"\"\n\n def __init__(self):\n super().__init__()\n self._attrs[\"op\"] = \"my_new_op\"\n\n def _infer_shape(self, x: Tensor) -> List[IntVar]:\n \"\"\"Compute output shape from input shapes.\"\"\"\n return x._attrs[\"shape\"] # adjust as needed\n\n def __call__(self, x: Tensor) -> Tensor:\n self._attrs[\"inputs\"] = [x]\n self._set_depth()\n output = Tensor(shape=self._infer_shape(x), src_ops={self})\n self._attrs[\"outputs\"] = [output]\n return output\n```\n\n#### 2. Write Backend Codegen\n\nAdd codegen templates in the appropriate backend directory:\n\n- **CUDA**: `python/aitemplate/backend/cuda/`\n- **ROCm**: `python/aitemplate/backend/rocm/`\n\nEach backend file typically provides:\n- `gen_function()` — generates the C++ kernel function\n- `gen_function_decl()` — generates the function declaration\n- `gen_function_call()` — generates the call site\n\n#### 3. Register the Operator\n\nMake sure the operator is importable from the ops package by adding it to the relevant `__init__.py`.\n\n#### 4. Add Tests\n\nCreate a test file in the appropriate test directory following existing patterns. Tests typically:\n- Build a small graph using the new op\n- Compile it with `compile_model()`\n- Compare outputs against a PyTorch reference implementation\n\n### Tips\n\n- Look at similar existing operators for patterns to follow\n- Gemm ops are in `python/aitemplate/compiler/ops/gemm/`\n- Elementwise ops are in `python/aitemplate/compiler/ops/common/`\n- Use `_attrs` dict to store operator metadata\n- Shape inference is critical — get it right or compilation will fail\n\n\n## AITemplate Architecture\n\n### Compilation Pipeline\n\nAITemplate transforms a high-level neural network graph into a compiled GPU binary through these stages:\n\n1. **Frontend (Python)**: Define or import a model using AIT operators (`python/aitemplate/compiler/ops/`)\n2. **Graph Optimization**: Fuse operators (horizontal, vertical, memory fusion) in `python/aitemplate/compiler/transform/`\n3. **Code Generation**: Emit C++ and CUDA/HIP source files via `python/aitemplate/backend/`\n4. **Compilation**: Invoke nvcc/hipcc to produce a `.so` shared library\n5. **Runtime**: Load the `.so` using the C++ runtime in `static/`\n\n### Key Modules\n\n- **`python/aitemplate/compiler/`**: Core compiler — IR, ops, graph transforms, type inference\n - `ops/`: Operator definitions (gemm, conv, elementwise, etc.)\n - `transform/`: Graph optimization passes (fusion, memory planning)\n - `base.py`: Base tensor and operator classes\n- **`python/aitemplate/backend/`**: Code generation backends\n - `cuda/`: NVIDIA-specific codegen using CUTLASS\n - `rocm/`: AMD-specific codegen using Composable Kernel\n - `common/`: Shared codegen utilities\n- **`python/aitemplate/frontend/`**: User-facing API for building models\n- **`python/aitemplate/utils/`**: Utility functions (shape inference, profiling, etc.)\n- **`static/`**: C++ runtime for loading and running compiled models\n\n### Operator Fusion Types\n\n- **Vertical fusion**: Fuse sequential ops (e.g., GEMM + bias + ReLU)\n- **Horizontal fusion**: Batch independent ops of the same type\n- **Memory fusion**: Eliminate intermediate allocations between fused ops\n\n### Backend Codegen\n\nEach backend (CUDA/ROCm) provides:\n- Kernel templates (Jinja2-based C++ templates)\n- Profiling to select optimal kernel configurations\n- Header generation for the compiled model API\n\n\n# Adding a CuTeDSL Backend to AITemplate: Execution Plan & Reusable Skill\n\nThis document captures the end-to-end process of converting an AITemplate (AIT)\noperator from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS\nPython DSL) with AOT compilation. It is generalized as a reusable pattern\napplicable to any AIT operator.\n\n---\n\n## Architecture Overview\n\n```\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CUTLASS C++ Backend (Original) │\n│ │\n│ Python (AIT codegen) C++ (generated) Binary │\n│ ┌─────────────────┐ ┌──────────────────┐ ┌──────────────────┐ │\n│ │ Jinja2 template │─────▶│ .cu │─nvcc──▶│ .obj │ │\n│ │ (FUNC_TEMPLATE) │ │ (CUTLASS C++ │ │ │ │\n│ │ │ │ template code) │ │ │ │\n│ └─────────────────┘ └──────────────────┘ └───────┬──────────┘ │\n│ │ link │\n│ ┌──────▼──────────┐ │\n│ │ test.so │ │\n│ └─────────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n\n┌───────────────────────────────────────────────────────────────────────────────┐\n│ CuTeDSL Backend (New) │\n│ │\n│ Python (AIT codegen + CuTeDSL AOT) C++ (thin wrapper) Binary │\n│ ┌─────────────────┐ ┌──────────────┐ ┌───────────────┐ ┌────────────┐ │\n│ │ CuTeDSL kernel │ │ cute.compile()│ │.cu │ │.obj │ │\n│ │ (@cute.jit + │──▶│ export_to_c()│──▶│(struct init + │──▶│(from nvcc) │ │\n│ │ @cute.kernel) │ │ │ │ metadata load)│ │ │ │\n│ └─────────────────┘ └──────┬───────┘ └───────────────┘ └────┬───────┘ │\n│ │ │ link │\n│ ┌─────▼────────┐ ┌─────▼───────┐ │\n│ │_cutedsl.h │ │ test.so │ │\n│ │_cutedsl.o │────────────────────▶│ (+libcuda) │ │\n│ │(pre-compiled) │ └─────────────┘ │\n│ └──────────────┘ │\n└───────────────────────────────────────────────────────────────────────────────┘\n```\n\n### Key Difference\n\n- **CUTLASS C++ backend**: Jinja2 template IS the kernel. It generates C++ source\n that instantiates CUTLASS template classes. `nvcc` compiles everything.\n- **CuTeDSL backend**: Kernel is written in Python using CuTeDSL decorators\n (`@cute.jit`, `@cute.kernel`). `cute.compile()` produces MLIR+cubin at\n codegen time, `export_to_c()` emits `.h`+`.o`. Only a thin C++ wrapper\n needs nvcc compilation.\n\n---\n\n## Step-by-Step Integration Guide\n\n### Step 1: Write the CuTeDSL Kernel\n\n**Create**: `backend/cuda//cutedsl__sm80.py`\n\nThis file contains the pure CuTeDSL kernel implementation:\n\n```python\nimport cutlass\nimport cutlass.cute as cute\nfrom cutlass.cute.nvgpu import cpasync, warp\n\nclass MyOpSm80Kernel:\n def __init__(self, ...compile_time_params...):\n self.param = param\n self.cta_sync_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=128)\n\n @cute.jit\n def __call__(self, mA: cute.Tensor, mB: cute.Tensor, ..., stream: cuda.CUstream):\n \"\"\"Host-side: create SMEM layouts, copy atoms, MMA config, launch kernel.\"\"\"\n # 1. Compute swizzled SMEM layouts\n sA_layout = cute.tile_to_shape(get_smem_layout_atom(...), ...)\n # 2. Create copy atoms (cpasync for GMEM->SMEM, universal for stores)\n gmem_copy = cute.make_tiled_copy_tv(cute.make_copy_atom(cpasync.CopyG2SOp(), ...), ...)\n # 3. Create tiled MMA\n tiled_mma = cute.make_tiled_mma(warp.MmaF16BF16Op(...), ...)\n # 4. Compute grid dimensions\n grid_dim = (m_blocks, batch_size, 1)\n # 5. Launch kernel\n self.kernel(...).launch(grid=grid_dim, block=[128, 1, 1], stream=stream)\n\n @cute.kernel\n def kernel(self, ...all_params...):\n \"\"\"Device-side: actual CUDA kernel body.\"\"\"\n tidx, _, _ = cute.arch.thread_idx()\n m_block, batch, _ = cute.arch.block_idx()\n smem = cutlass.utils.SmemAllocator()\n # ... GEMM mainloop, epilogue, store ...\n```\n\n**Key patterns**:\n\n1. **Swizzled SMEM layouts** (for bank-conflict-free access):\n ```python\n def get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = (128 if bytes_per_row % 128 == 0 else ...) // dtype_byte\n swizzle_bits = 4 if smem_k_block_size == 128 else ...\n return cute.make_composed_layout(\n cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), 0,\n cute.make_ordered_layout((8, smem_k_block_size), order=(1, 0)),\n )\n ```\n\n2. **SM80 warp-level MMA** (`warp.MmaF16BF16Op`):\n ```python\n tiled_mma = cute.make_tiled_mma(\n warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)),\n (num_warps, 1, 1),\n permutation_mnk=(num_warps * 16, 16, 16),\n )\n ```\n\n3. **cp.async for GMEM->SMEM**:\n ```python\n atom_async = cute.make_copy_atom(\n cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),\n dtype, num_bits_per_copy=128)\n gmem_tiled_copy = cute.make_tiled_copy_tv(atom_async, thread_layout, value_layout)\n ```\n\n4. **SMEM->register via ldmatrix**:\n ```python\n smem_copy_atom = cute.make_copy_atom(\n warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), dtype)\n smem_tiled_copy_A = cute.make_tiled_copy_A(smem_copy_atom, tiled_mma)\n ```\n\n5. **Activation dispatch table**:\n ```python\n ACTIVATION_MAP = {\n \"Sigmoid\": lambda x: 1.0 / (1.0 + cute.math.exp(-x)),\n \"ReLu\": lambda x: cute.arch.fmax(x, 0.0),\n \"Identity\": lambda x: x,\n }\n ```\n\n### Step 2: Create the Backend Codegen Module\n\n**Create**: `backend/cuda//_cutedsl.py`\n\nThis file registers AIT backend functions with `_cutedsl` suffix:\n\n```python\nimport functools, os, jinja2\nfrom aitemplate.backend import registry\nfrom aitemplate.backend.target import Target\n\n# C++ wrapper template - constructs CuTeDSL typed tensor structs from AIT void* pointers\nCUTEDSL_WRAPPER_TEMPLATE = jinja2.Template(\"\"\"\n#include \n#include \"{{cutedsl_header}}\"\n\nstatic {{func_name}}_cutedsl_Kernel_Metadata_t g_metadata;\nstatic bool g_metadata_loaded = false;\n\nstatic void ensure_metadata_loaded() {\n if (!g_metadata_loaded) {\n cuInit(0);\n {{func_name}}_cutedsl_Kernel_Metadata_Load(&g_metadata);\n g_metadata_loaded = true;\n }\n}\n\nvoid {{func_name}}(...AIT signature...) {\n ensure_metadata_loaded();\n\n // Construct CuTeDSL typed tensor structs:\n {{func_name}}_cutedsl_Tensor_mA_t t_mA;\n t_mA.data = a_ptr;\n t_mA.dynamic_shapes[0] = batch; // first dynamic dim\n t_mA.dynamic_shapes[1] = seq_len; // second dynamic dim\n t_mA.dynamic_strides[0] = seq_len * K; // batch stride\n\n // ... same for all tensors ...\n\n CUstream cu_stream = reinterpret_cast(stream);\n {{func_name}}_cutedsl_wrapper(&g_metadata, &t_mA, ..., cu_stream);\n}\n\"\"\")\n\n@functools.lru_cache(maxsize=32)\ndef _aot_compile_cutedsl_kernel(n0, n1, ..., output_dir, func_name, arch):\n \"\"\"AOT compile kernel and export .h + .o.\"\"\"\n from my_kernel import MySm80Kernel\n from cutlass.cute.runtime import from_dlpack\n from cutlass.cute.export import export_to_c\n\n kernel = MySm80Kernel(n0=n0, n1=n1, ...)\n\n # Create representative tensors (shapes don't matter for dynamic dims)\n q_pt = torch.zeros(4, 128, n1, device=\"cuda\", dtype=torch.float16)\n\n # Mark dynamic modes with per-mode divisibility\n q_cute = from_dlpack(q_pt, assumed_align=16)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=0, stride_order=q_pt.dim_order(), divisibility=1)\n q_cute = q_cute.mark_compact_shape_dynamic(\n mode=1, stride_order=q_pt.dim_order(), divisibility=8)\n\n # Compile and export\n compiled = cute.compile(kernel, q_cute, ..., cu_stream)\n export_to_c(compiled, file_path=output_dir, file_name=f\"{func_name}_cutedsl\")\n return h_path, o_path\n\n@registry.reg(\"cuda..gen_function_cutedsl\")\ndef gen_function_cutedsl(func_attrs):\n h_path, o_path = _aot_compile_cutedsl_kernel(...)\n func_attrs[\"cutedsl_obj_path\"] = o_path # For builder to pick up\n return CUTEDSL_WRAPPER_TEMPLATE.render(...)\n\n@registry.reg(\"cuda..func_decl_cutedsl\")\ndef func_decl_cutedsl(func_attrs): ...\n\n@registry.reg(\"cuda..func_call_cutedsl\")\ndef func_call_cutedsl(func_attrs, indent=\" \"): ...\n```\n\n**Critical detail**: The C++ wrapper MUST construct CuTeDSL typed tensor structs\n(`*_Tensor_mX_t`) with proper `dynamic_shapes[]` and `dynamic_strides[]` fields.\nThe `export_to_c()` API generates these structs -- they are NOT raw `void*`.\n\n### Step 3: Add Backend Dispatch in the Op Class\n\n**Modify**: `compiler/ops//.py`\n\nAdd methods to detect CuTeDSL mode and dispatch with suffix:\n\n```python\ndef _use_cutedsl(self) -> bool:\n current_target = target.Target.current()\n return current_target._kwargs.get(\"use_cutedsl_\", False)\n\ndef _backend_suffix(self) -> str:\n return \"_cutedsl\" if self._use_cutedsl() else \"\"\n\ndef gen_function(self) -> str:\n suffix = self._backend_suffix()\n self._attrs[\"backend_suffix\"] = suffix # For codegen.py to pick up\n func_key = f\"{target_name}.{op_name}.gen_function{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_decl(self, func_attrs=None) -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_decl{suffix}\"\n return registry.get(func_key)(self._attrs)\n\ndef gen_function_call(self, func_attrs=None, indent=\" \") -> str:\n suffix = self._backend_suffix()\n func_key = f\"{target_name}.{op_name}.func_call{suffix}\"\n return registry.get(func_key)(self._attrs, indent)\n```\n\n### Step 4: Modify the Build System\n\n#### 4.1 `backend/codegen.py` -- `gen_function_src()`\n\nPass `workdir` to `func_attrs` so the AOT compiler knows where to write:\n```python\nfunc._attrs[\"workdir\"] = prefix\n```\n\nAfter writing the `.cu` wrapper, add the pre-compiled `.o` to `file_pairs`:\n```python\ncutedsl_obj = func._attrs.get(\"cutedsl_obj_path\")\nif cutedsl_obj and os.path.exists(cutedsl_obj):\n file_pairs.append((cutedsl_obj, cutedsl_obj))\n```\n\nIn `_process_src_ops()`, read `backend_suffix` for `func_decl` / `func_call` dispatch:\n```python\nbackend_suffix = func._attrs.get(\"backend_suffix\", \"\")\nf_func_decl = registry.get(f\"{target}.{op}.func_decl{backend_suffix}\")\nf_func_call = registry.get(f\"{target}.{op}.func_call{backend_suffix}\")\n```\n\n#### 4.2 `backend/builder.py` -- Handle pre-compiled `.o` files\n\nIn `build_objs()`: skip compilation for `.o` source files:\n```python\nif src.endswith(\".o\"):\n _LOGGER.info(f\"Skipping compilation for pre-compiled object: {src}\")\n continue\n```\n\nIn `build_so()`: add `-lcuda` when CuTeDSL objects are present:\n```python\nhas_cutedsl = any(obj.endswith(\"_cutedsl.o\") for obj in objs)\nextra_libs = \" -lcuda\" if has_cutedsl else \"\"\n```\n\nIn `gen_makefile()`: use `os.path.basename()` for `.o` paths (Make runs from workdir):\n```python\nif pair[0].endswith(\".o\"):\n obj_files.append(os.path.basename(pair[0])) # NOT full path\n has_cutedsl = True\n```\n\n### Step 5: Update Target Configuration\n\n**Modify**: `backend/target.py`\n\nAdd `use_cutedsl_` kwarg support to the Target class (already passed through\n`_kwargs`; no explicit changes needed unless you want validation).\n\n### Step 6: Update `__init__.py` and Examples\n\n**Modify**: `backend/cuda//__init__.py`\n\nAdd the new module import for auto-registration:\n```python\nfrom aitemplate.backend.cuda. import _cutedsl\n```\n\n**Modify**: Example/test scripts\n\nAdd CLI arg to enable CuTeDSL:\n```python\ntarget = Target(..., use_cutedsl_=True)\n```\n\n### Step 7: Add BUCK Dependencies\n\nAdd `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps\nin `examples/BUCK` for any target that uses the CuTeDSL backend.\n\n---\n\n## Common Pitfalls and Debug Checklist\n\n| # | Error | Root Cause | Fix |\n|---|-------|------------|-----|\n| 1 | `ModuleNotFoundError: No module named 'cuda'` | Missing `nvidia-cutlass-dsl` BUCK dep | Add `fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl` to deps |\n| 2 | `ValueError: Invalid leading dimension: 2` | `LayoutEnum.from_tensor()` only accepts 2D tensors (leading_dim 0 or 1) | Hardcode `utils.LayoutEnum.ROW_MAJOR` for 3D row-major tensors |\n| 3 | `_ScaledBasis TypeError` in SM90 TMA | CuTeDSL TMA layout algebra doesn't yet handle 3D tensors well | Use SM80 kernel (works on SM90+ hardware too) |\n| 4 | `coord and shape of view are weakly congruent: N modes vs M coords` | After `partition_S(local_tile(..., (tile, None)))`, the tile mode creates an extra dimension | Add extra `None` to indexing: `tQgQ[None, None, None, k_tile]` instead of `tQgQ[None, None, k_tile]` |\n| 5 | `ptr alignment does not meet requirement` for 128-bit cp.async | `mark_layout_dynamic()` loses alignment info | Use `mark_compact_shape_dynamic(mode=M, divisibility=8)` to preserve alignment |\n| 6 | `cute.copy expects src and dst to have static shape` | `mark_layout_dynamic()` makes ALL dims dynamic, but inner dims must be static constants | Use per-mode `mark_compact_shape_dynamic()` only on truly dynamic modes (batch, seq_len); leave inner dims (N0, N1, K) static |\n| 7 | `shape(X) of mode(0) is not divisible by divisibility(8)` | Batch dim size not divisible by the divisibility requirement | Use per-mode divisibility: batch `div=1`, M/seq_len `div=8` |\n| 8 | `argument of type \"void *\" incompatible with \"..._Tensor_mX_t *\"` | CuTeDSL `export_to_c` generates typed tensor structs, not raw pointer APIs | Construct typed tensor structs in C++ wrapper: set `.data`, `.dynamic_shapes[]`, `.dynamic_strides[]` |\n| 9 | `No rule to make target 'tmp/.../_cutedsl.o'` | Makefile has full path but Make runs from workdir | Use `os.path.basename(pair[0])` in `gen_makefile()` |\n| 10 | `DT_TEXTREL` linker warning | CuTeDSL `.o` contains text relocations | Non-fatal warning; can be ignored |\n\n---\n\n## Dynamic Tensor Marking Strategy\n\nThis is the most critical part for correctness. CuTeDSL needs to know which\ntensor dimensions are compile-time constants vs runtime-dynamic.\n\n### DO NOT use `mark_layout_dynamic()`\n\n`mark_layout_dynamic(leading_dim=N)` makes ALL shapes and strides dynamic.\nThis breaks:\n- `cute.copy` which requires static shapes for src/dst\n- cp.async alignment checks (loses divisibility info)\n- Optimal code generation (no compile-time constants)\n\n### DO use `mark_compact_shape_dynamic()` per mode\n\n```python\ndef make_cute_tensor(torch_tensor, dynamic_modes_div):\n \"\"\"\n dynamic_modes_div: list of (mode_index, divisibility) tuples.\n Only the specified modes become dynamic; all other dims stay static.\n \"\"\"\n ct = from_dlpack(torch_tensor, assumed_align=16)\n for mode, div in dynamic_modes_div:\n ct = ct.mark_compact_shape_dynamic(\n mode=mode,\n stride_order=torch_tensor.dim_order(),\n divisibility=div,\n )\n return ct\n```\n\n**Typical b2b_bmm pattern**:\n| Tensor | Shape | Dynamic Modes |\n|--------|-------|---------------|\n| Q | [batch, M, K] | batch (div=1), M (div=8) |\n| K | [batch, N0, K] | batch (div=1) only |\n| V | [batch, N0, N1] | batch (div=1) only |\n| Bias | [batch, M, N0] | batch (div=1), M (div=8) |\n| Out | [batch, M, N1] | batch (div=1), M (div=8) |\n\n### Why div=8 for M?\n\n128-bit cp.async copies 8 fp16 elements at once. The M dimension appears in\nstrides that feed cp.async, so CuTeDSL needs `divisibility=8` to guarantee\nalignment for 128-bit vectorized copies.\n\n---\n\n## Generated Artifact Reference\n\nAfter successful compilation, the workdir contains:\n\n| File | Size (typical) | Description |\n|------|----------------|-------------|\n| `_cutedsl.h` | ~4 KB | CuTeDSL header: typed tensor structs, `Metadata_Load/Unload`, `wrapper()` |\n| `_cutedsl.o` | ~50 KB | Pre-compiled object with embedded cubin(s) |\n| `.cu` | ~4 KB | Thin C++ wrapper: `cuInit`, metadata load, struct construction, launch |\n| `.obj` | ~16 KB | Compiled wrapper object |\n| `test.so` | ~1.4 MB | Final shared library (linked with `-lcuda`) |\n\n### Header structure (`_cutedsl.h`)\n\n```c\n// Metadata struct (holds CUlibrary handle)\ntypedef struct {\n CUlibrary module;\n} _cutedsl_Kernel_Metadata_t;\n\n// Typed tensor structs (one per kernel argument)\ntypedef struct {\n void *data;\n int32_t dynamic_shapes[N_DYNAMIC_SHAPES]; // runtime batch, seq_len, etc.\n int64_t dynamic_strides[N_DYNAMIC_STRIDES]; // runtime strides\n} _cutedsl_Tensor_mX_t;\n\n// Load/Unload/Wrapper functions\nvoid _cutedsl_Kernel_Metadata_Load(_cutedsl_Kernel_Metadata_t *metadata);\nvoid _cutedsl_Kernel_Metadata_Unload(_cutedsl_Kernel_Metadata_t *metadata);\nint32_t _cutedsl_wrapper(_cutedsl_Kernel_Metadata_t *metadata,\n _cutedsl_Tensor_mX_t *mX, ..., CUstream stream);\n```\n\n---\n\n## File Checklist (for classic_b2b_bmm)\n\n### Files Created\n\n| File | Path (relative to aitemplate/) |\n|------|-------------------------------|\n| SM80 kernel | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` |\n| SM90 kernel (WIP) | `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm90.py` |\n| Backend codegen | `backend/cuda/b2b_bmm/classic_b2b_bmm_cutedsl.py` |\n| AOT example | `examples/07_how_to_run_pt_model/cutedsl_b2b_bmm_aot_example.py` |\n\n### Files Modified\n\n| File | Change |\n|------|--------|\n| `compiler/ops/b2b_bmm/classic_b2b_bmm.py` | Added `_use_cutedsl()`, `_backend_suffix()`, dispatch in `gen_function/decl/call` |\n| `backend/codegen.py` | Pass `workdir`, handle `.o` in `file_pairs`, read `backend_suffix` |\n| `backend/builder.py` | Skip `.o` compilation, add `-lcuda`, use `basename()` in Makefile |\n| `backend/cuda/b2b_bmm/__init__.py` | Added `classic_b2b_bmm_cutedsl` import |\n| `examples/07_how_to_run_pt_model/classic_b2b_bmm_example.py` | Added `--use-cutedsl` and `--both` CLI args |\n| `examples/BUCK` | Added `nvidia-cutlass-dsl` dependency |\n\n---\n\n## Reference Codebases\n\n| Reference | Path | What to reuse |\n|-----------|------|---------------|\n| FA4 SM80 helpers | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | SM80 warp MMA, SMEM layout, gemm/gemm_rs patterns |\n| FA4 SM90 forward | `fbcode/ads_mkl/ops/cute_dsl/fa4/src/flash_fwd.py` | SM90 TMA+WGMMA, pipeline patterns |\n| Hopper GEMM example | `fbcode/ai_acceleration/cute_dsl/examples/hopper/dense_gemm.py` | Complete CuTeDSL GEMM lifecycle |\n| CuTeDSL AOT export | `third-party/cutlass/4.3.5/python/CuTeDSL/cutlass/cute/export/` | `export_to_c()`, `dump_to_object()` API |\n| Quack PT2 compat | `fbcode/ads_mkl/ops/cute_dsl/quack/quack_rmsnorm_pt2.py` | `@functools.cache` + `cute.compile()` caching |\n| Inductor CuTeDSL | `fbcode/caffe2/torch/_inductor/codegen/cutedsl/` | Framework integration pattern |\n| Current CUTLASS backend | `backend/cuda/b2b_bmm/classic_b2b_bmm.py` | Interface contract (signature, registry keys) |\n| CUTLASS C++ kernels | `static/include/kernels/classic_b2b_bmm/` | Algorithm reference for register-passing b2b GEMM |\n\n---\n\n## Claude Skill: `add-cutedsl-backend`\n\nBelow is a reusable Claude skill specification for adding a CuTeDSL backend\nto any AIT operator.\n\n### Skill Description\n\n```\nWhen adding a CuTeDSL backend for an AITemplate operator, follow these steps:\n\n1. RESEARCH: Read the existing CUTLASS C++ backend for the target op:\n - `backend/cuda//.py` -- understand FUNC_TEMPLATE, registry keys,\n tensor shapes, tiling parameters\n - `compiler/ops//.py` -- understand gen_function/decl/call dispatch\n\n2. KERNEL: Write the CuTeDSL kernel in `backend/cuda//cutedsl__sm80.py`:\n - Class with __init__ (compile-time params), @cute.jit __call__ (host setup),\n @cute.kernel kernel (device code)\n - SM80: use warp.MmaF16BF16Op, cpasync.CopyG2SOp, warp.LdMatrix8x8x16bOp\n - Swizzled SMEM layouts via get_smem_layout_atom()\n - Named barriers for CTA sync\n\n3. CODEGEN: Create `backend/cuda//_cutedsl.py`:\n - @functools.lru_cache AOT compile function\n - C++ wrapper template with typed tensor struct construction\n - Registry functions with _cutedsl suffix\n - CRITICAL: Use mark_compact_shape_dynamic() per-mode, NOT mark_layout_dynamic()\n - CRITICAL: C++ wrapper must construct *_Tensor_mX_t structs, NOT pass void*\n\n4. DISPATCH: Modify `compiler/ops//.py`:\n - Add _use_cutedsl(), _backend_suffix()\n - Add suffix dispatch in gen_function(), gen_function_decl(), gen_function_call()\n - Store backend_suffix in self._attrs\n\n5. BUILD: Modify `backend/codegen.py` and `backend/builder.py`:\n - codegen.py: pass workdir, handle .o in file_pairs, read backend_suffix\n - builder.py: skip .o compilation, add -lcuda, use basename() for .o in Makefile\n\n6. REGISTER: Add import in `backend/cuda//__init__.py`\n\n7. TEST: Run with Target(..., use_cutedsl_=True) and validate against\n PyTorch reference (atol=1e-2 for fp16)\n\nCommon pitfalls to watch for:\n- mark_layout_dynamic makes ALL dims dynamic -> breaks static copy requirements\n- partition_S(local_tile(..., (tile, None))) adds extra modes -> count None indices\n- CuTeDSL export_to_c generates typed structs, not raw void* -> read .h header\n- Makefile .o paths must use basename (Make runs from workdir)\n- Need -lcuda for CUlibrary-based metadata loading\n```\n\n---\n\n## Execution Order for New Operators\n\n1. **Read existing backend** -- understand tensor shapes, tiling, registry keys\n2. **Write SM80 kernel** -- implement, test standalone with `cute.compile()`\n3. **AOT export test** -- verify `export_to_c()` produces valid `.h` + `.o`\n4. **Create backend codegen** -- registry functions, C++ wrapper template\n5. **Add dispatch** -- op class modifications, `__init__.py` import\n6. **Modify build system** -- codegen.py, builder.py changes\n7. **End-to-end test** -- run full AIT pipeline, compare outputs\n8. **(Optional) SM90 kernel** -- TMA + WGMMA version for Hopper\n\n---\n\n## Verification Checklist\n\n- [ ] CuTeDSL kernel compiles: `cute.compile()` succeeds\n- [ ] AOT export produces `.h` and `.o` files\n- [ ] C++ wrapper compiles with nvcc (no type mismatches)\n- [ ] `.o` links into `.so` (no undefined symbols, `-lcuda` present)\n- [ ] Runtime loads metadata: `Kernel_Metadata_Load()` succeeds\n- [ ] Numerical correctness: output matches PyTorch reference (atol<=0.02 for fp16)\n- [ ] Dynamic shapes work: different batch sizes / sequence lengths\n- [ ] Original CUTLASS backend still works (no regression)\n\n\n## Migrating GEMM Operators from CUTLASS C++ to CuTeDSL\n\n### Overview\n\nThis skill documents the pattern for migrating any GEMM variant (gemm_rcr_bias, gemm_rrr, gemm_ccr, etc.) from CUTLASS C++ template-based code generation to CuTeDSL (CUTLASS Python DSL), using AOT-style deployment.\n\n### Layout Reference\n\nCUTLASS GEMM naming convention: `gemm_` where R=Row-major, C=Column-major.\n\n| Layout | A shape/order | B shape/order | MMA Atom | Key property |\n|--------|--------------|--------------|----------|-------------|\n| RCR | [M,K] K-major | [N,K] K-major | TN | Both K-contiguous |\n| RRR | [M,K] K-major | [K,N] N-major | TN/NT | A=K-cont, B=N-cont |\n| CCR | [K,M] M-major | [N,K] K-major | NN | A=M-cont, B=K-cont |\n\nFor SM80 with fp16: `warp.MmaF16BF16Op(Float16, Float32, (16,8,16))` with TN layout.\n\n### Step-by-Step Migration\n\n#### 1. Identify the operation\n\nRead the CUTLASS backend file (e.g., `backend/cuda/gemm_universal/gemm_rcr_bias.py`) to understand:\n- The GEMM layout (RCR, RRR, etc.)\n- The epilogue (bias, relu, bias+relu, etc.)\n- The function signature from `common_bias.py` or `common.py`\n- Whether it supports split-K\n\n#### 2. Create the CuTeDSL kernel\n\nCreate `backend/cuda/gemm_universal/cutedsl__sm80.py`.\n\nKey components to implement:\n\n```python\nclass GemmXxxSm80Kernel:\n def __init__(self, tile_m=128, tile_n=128, tile_k=32):\n self.tile_m, self.tile_n, self.tile_k = tile_m, tile_n, tile_k\n self.num_threads = 128 # 4 warps\n\n @cute.jit\n def __call__(self, mA, mB, mC, ...):\n # 1. Set up SMEM layouts using get_smem_layout_atom()\n # 2. Define SharedStorage struct\n # 3. Create copy atoms (async for load, universal for store)\n # 4. Create tiled_mma with SM80 16x8x16 atom\n # 5. Compute grid dims and launch kernel\n\n @cute.kernel\n def kernel(self, ...):\n # 1. Allocate SMEM, partition GMEM tiles\n # 2. GEMM mainloop: K-loop with cp.async + warp MMA\n # 3. Epilogue: apply bias/activation on FP32 accumulators\n # 4. Store: Reg -> SMEM -> GMEM\n```\n\n**SMEM layout pattern** (from `cutedsl_b2b_bmm_sm80.py`):\n```python\ndef get_smem_layout_atom(dtype, k_dim):\n dtype_byte = cutlass.const_expr(dtype.width // 8)\n bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)\n smem_k_block_size = cutlass.const_expr(\n 128 if bytes_per_row % 128 == 0\n else (64 if bytes_per_row % 64 == 0\n else (32 if bytes_per_row % 32 == 0 else 16))\n ) // dtype_byte\n # ... swizzle computation\n return cute.make_composed_layout(swizzle, 0, ordered_layout)\n```\n\n**GEMM mainloop pattern**:\n```python\nfor k_tile in range(k_tiles):\n cute.copy(gmem_tiled_copy, src_gmem, dst_smem)\n cute.arch.cp_async_commit_group()\n cute.arch.cp_async_wait_group(0)\n barrier.arrive_and_wait()\n\n # Load first k-block from SMEM to registers\n cute.copy(smem_tiled_copy_A, tSsA[None, None, 0], tSrA_view[None, None, 0])\n cute.copy(smem_tiled_copy_B, tSsB[None, None, 0], tSrB_view[None, None, 0])\n\n for k_block in cutlass.range_constexpr(num_k_blocks):\n k_next = (k_block + 1) % num_k_blocks\n cute.copy(smem_tiled_copy_A, tSsA[None, None, k_next], ...)\n cute.copy(smem_tiled_copy_B, tSsB[None, None, k_next], ...)\n cute.gemm(tiled_mma, acc, tSrA[None, None, k_block], tSrB[None, None, k_block], acc)\n\n barrier.arrive_and_wait()\n```\n\n**Store pattern** (Reg -> SMEM -> GMEM):\n```python\nrOut = cute.make_fragment_like(acc, dtype)\nrOut.store(acc.load().to(dtype))\nsOut = cute.make_tensor(sA.iterator, sOut_layout) # reuse SMEM\n\nsmem_copy_atom_C = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype)\nsmem_tiled_copy_C = cute.make_tiled_copy_C(smem_copy_atom_C, tiled_mma)\n# ... copy from reg to SMEM, sync, then SMEM to GMEM\n```\n\n#### 3. Handle the epilogue\n\n**1D bias (broadcast)**: Load bias into SMEM as a 2D view with stride-0 in M, use `thr_mma.partition_C()` to align with accumulator layout:\n```python\n# In __call__: create 2D broadcast view\nmBias_2d = cute.make_tensor(mBias.iterator,\n cute.make_layout((M_size, N_size), stride=(0, bias_stride)))\n\n# In kernel: tile and load\ngBias = cute.local_tile(mBias_2d, (tile_m, tile_n), (m_block, n_block))\n# ... cp.async load to sBias ...\ntCsBias = thr_mma.partition_C(sBias)\n# element-wise add in FP32\n```\n\n**2D bias**: Direct tile and load (like b2b_bmm does).\n\n**Activation**: Apply element-wise in FP32 accumulator space before store.\n\n#### 4. Create AOT example\n\nCreate `examples/07_how_to_run_pt_model/cutedsl__aot_example.py`:\n\n```python\ndef make_cute_tensor(torch_tensor):\n ct = from_dlpack(torch_tensor, assumed_align=16)\n ct = ct.mark_layout_dynamic(leading_dim=len(torch_tensor.shape) - 1)\n ct = ct.mark_compact_shape_dynamic(\n mode=len(torch_tensor.shape) - 1,\n stride_order=torch_tensor.dim_order(),\n divisibility=(128 // cutlass.Float16.width),\n )\n return ct\n\ndef jit_compile_and_validate():\n # 1. Create PyTorch reference tensors\n # 2. Compute PyTorch reference output\n # 3. Create kernel instance\n # 4. cute.compile(kernel, *cute_tensors, *scalar_args, stream)\n # 5. Execute and compare (atol=1e-2, rtol=1e-2 for fp16)\n\ndef aot_export(compiled):\n # export_to_c(compiled, file_path=output_dir, file_name=name)\n # Generates: .h + .o\n```\n\n#### 5. Add BUCK target\n\n```python\npython_binary(\n # @autodeps-skip\n name = \"cutedsl__aot_example\",\n srcs = [\"07_how_to_run_pt_model/cutedsl__aot_example.py\"],\n main_module = \"aitemplate.AITemplate.examples.07_how_to_run_pt_model.cutedsl__aot_example\",\n deps = [\n \"fbsource//third-party/pypi/nvidia-cutlass-dsl:nvidia-cutlass-dsl\",\n \"//aitemplate/AITemplate/python/aitemplate:aitemplate\",\n \"//caffe2:torch\",\n ],\n)\n```\n\n### Reference Implementations\n\n| File | What it demonstrates |\n|------|---------------------|\n| `backend/cuda/gemm_universal/cutedsl_gemm_rcr_bias_sm80.py` | GEMM RCR + 1D bias, 2D grid |\n| `backend/cuda/b2b_bmm/cutedsl_b2b_bmm_sm80.py` | Fused B2B GEMM, 2D bias, activation |\n| `examples/.../cutedsl_gemm_rcr_bias_aot_example.py` | AOT example for GEMM |\n| `examples/.../cutedsl_b2b_bmm_aot_example.py` | AOT example for B2B BMM |\n| `ads_mkl/ops/cute_dsl/fa4/src/ampere_helpers.py` | Reusable SM80 helpers |\n\n### Common Pitfalls\n\n1. **Alignment**: Use `mark_compact_shape_dynamic(divisibility=8)` for fp16 to satisfy 128-bit cp.async alignment\n2. **Bias broadcast**: Create 2D view with stride-0 in broadcast dim at host level (`__call__`), not inside `@cute.kernel`\n3. **SMEM reuse**: Output staging can reuse SMEM from input tiles (e.g., `sC = cute.make_tensor(sA.iterator, sC_layout)`)\n4. **MN view**: Use `_make_acc_tensor_mn_view()` to iterate accumulator in logical (M, N) order for epilogue ops\n5. **Python slicing**: CuTe tensors don't support Python slicing (`tensor[start:]`). Use `cute.local_tile()` or `cute.make_tensor()` with offset iterators\n\n\n## Testing in AITemplate\n\n### Running Tests\n\n#### Python tests with pytest\n```bash\n# Run a specific test file\npython -m pytest tests/unittest/test_my_op.py -v\n\n# Run a specific test case\npython -m pytest tests/unittest/test_my_op.py::TestMyOp::test_basic -v\n```\n\n#### Buck tests (Meta-internal)\n```bash\n# Run a specific Buck test target\nbuck2 test //aitemplate/AITemplate:test_target_name\n```\n\n### Test Patterns\n\nAITemplate tests typically follow this structure:\n\n```python\nimport unittest\nimport torch\nfrom aitemplate.testing import detect_target\nfrom aitemplate.compiler import compile_model\nfrom aitemplate.frontend import Tensor, nn\n\nclass TestMyFeature(unittest.TestCase):\n def test_basic(self):\n # 1. Check for available GPU target\n target = detect_target()\n\n # 2. Build AIT graph\n X = Tensor(shape=[2, 3], name=\"X\", is_input=True)\n Y = my_op()(X)\n Y._attrs[\"is_output\"] = True\n Y._attrs[\"name\"] = \"Y\"\n\n # 3. Compile\n module = compile_model(Y, target, \"./tmp\", \"test_my_op\")\n\n # 4. Run with real data\n x_pt = torch.randn(2, 3).cuda().half()\n y = torch.empty(2, 3).cuda().half()\n module.run_with_tensors([x_pt], [y])\n\n # 5. Compare against PyTorch reference\n y_ref = torch_reference_impl(x_pt)\n torch.testing.assert_close(y, y_ref)\n```\n\n### Key Testing Utilities\n\n- `aitemplate.testing.detect_target()` — auto-detect CUDA or ROCm\n- `compile_model()` — compile a graph to a loadable module\n- `module.run_with_tensors()` — execute the compiled model\n- Tests typically use `float16` precision and compare with tolerance\n\n### FX2AIT Tests\n\nFor the PyTorch FX converter, tests are in `fx2ait/fx2ait/test/`:\n```python\n# These tests verify that PyTorch models convert correctly to AIT\n# and produce numerically equivalent results\n```\n"}