Add CuTe DSL Python ports for 01-14 (TVM-FFI enabled) - #16
Merged
Conversation
Each example in the series now ships a ``cutedsl_*.py`` port alongside
the original ``.cu`` / ``.py`` pair. The DSL ports skip the C++ build
step entirely — no template compilation, no torch.utils.cpp_extension
load — just an ``import`` of nvidia-cutlass-dsl and a one-time
``cute.compile(...)`` per dtype / is_gemm specialization.
Coverage (13 ports, 10 is skipped — no DSL analogue of the CUTLASS
device-level GemmUniversalAdapter / CollectiveBuilder API):
* 01 minimal-gemm — single 16x8x8 SM80 atom
* 02 mixed-precision-gemm — fp16/bf16 + per-dtype MMA op selection
* 03 tiled-mma — TiledMma with atom_layout + permutation
* 04 tiled-copy — TiledCopy via make_tiled_copy_tv
* 05 block-mma — block-grid TiledMma
* 06 block-copy — full G2S/S2R/R2S/S2G chain with swizzled
smem-staged epilogue
* 07 swizzling — Swizzle<3,3,3> smem (3-D tile_to_shape +
ComposedLayout)
* 08 dynamic-mma — predicated G2S + K-tile loop + R2S->S2G
epilogue, three dtype specs
* 09 pipelining — 3-stage cp.async pipeline (tensorop_gemm.py
mainloop pattern) + reg prefetch
* 11 tma-load-store — SM90 TMA bulk-tensor loads + R2S->TMA
store via PipelineTmaAsync
* 12 tma-multicast-reduce — cluster(2,1) + TMA-multicast on A +
TMA_REDUCE_ADD for the output
* 13 warpgroup-mma — Hopper WGMMA (OperandSource.SMEM, K_SW128)
* 14 warp-specialization — producer/consumer warpgroups, named
barrier, register reconfig
TVM-FFI plumbing (env-stream pattern, per the user docs):
* ``from_dlpack(t, assumed_align=16, enable_tvm_ffi=True)`` inside
``make_cute_tensor`` (kept as a thin helper, used for compile
templates only).
* ``cute.compile(..., options="--enable-tvm-ffi")``.
* Each ``@cute.jit`` host fn takes ``stream: CUstream`` and forwards
it to ``.launch(stream=stream)``.
* Compile passes ``make_fake_stream(use_tvm_ffi_env_stream=True)``;
runtime call sites pass raw ``torch.Tensor`` directly (no make_
cute_tensor wrap, no explicit stream — DSL syncs to
``torch.cuda.current_stream()`` for us).
C++ fidelity rule: every cutedsl mirrors its .cu kernel structure
(TiledMMA / TiledCopy shapes, smem swizzle, pipeline staging, TMA
multicast / reduce-add). Two intentional deviations (both flagged in
the file headers): (a) 08/09 first-tile G2S A/B uses a per-element
(M-bound ∧ K-bound) pred rather than the C++ iter-level gate, because
the gate silently drops valid threads when K < BLK_K — the C++
harness would itself fail tiny-K shapes; (b) 11/13/14 keep sC/sD as
separate smem buffers rather than aliasing on sA/sB — the H200 228KB
dynamic smem ceiling fits both layouts comfortably.
Sweeps match the original ``*.py`` harnesses exactly (08's N list also
gets ``64`` to align with M/K — dynamic_mma.py updated accordingly).
README documents the install requirements (nvidia-cutlass-dsl,
cuda-python + cuda-bindings, apache-tvm-ffi, torch_c_dlpack_ext) and
the env-stream invocation idiom.
Verified on G2 (H200, SM90) across the full sweeps:
01:2 | 02:4 | 03:2 | 04:2 | 05:2 | 06:2 | 07:2 | 08:1372 |
09:2048 | 11:2048 | 12:640 | 13:2048 | 14:2048 — all Failed=0
(10,316 / 10,316 total).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* ``third-party/cutlass`` advanced to the v4.5.0 release tag — gives us the canonical Hopper DSL examples (dense_gemm.py, tensorop_gemm.py) that 07-14's CuTe DSL ports cribbed from. * ``code-style-check.yml`` no longer dispatches the auto-fix workflow on failure (style failures stay as failed checks; humans rerun ``make style`` locally). * ``auto-fix-code-style.yml`` removed entirely. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
clang-format-21 reflowed a single ``using Spec = ...`` template parameter list in each of pipelining.cu, pipelining_no_reg_prefetch.cu, and gemm_api.cu. Pure whitespace; no semantic change. Brings ``make cquality`` to a clean state. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
cutedsl_*.pyPython ports alongside each example's existing.cu/.pypair, covering 01–09 + 11–14 (10-gemm-api skipped — no DSL analogue of CUTLASS'sGemmUniversalAdapter/CollectiveBuilderAPI).options="--enable-tvm-ffi"+make_fake_stream(use_tvm_ffi_env_stream=True); runtime call sites take baretorch.Tensorand ridetorch.cuda.current_stream()automatically.cutlasssubmodule bumped to v4.5.0 (gives us the canonicaldense_gemm.py/tensorop_gemm.pyDSL references); the auto-fix-code-style CI workflow is removed (style failures no longer triggers a dispatched fix).Install requirements (new in README):