From 1e7394829291360bdcf07036cbe5411631d2d33b Mon Sep 17 00:00:00 2001 From: Junkai-Wu Date: Tue, 14 Jul 2026 10:49:52 +0800 Subject: [PATCH] v4.6.1 update. (#3383) * v4.6.1 update. * Correct version number. * Fix documentation link for GEMM performance measurement Updated the link in the documentation release note for GEMM performance measurement methodology. * Fix URL in 'What's New in CUTLASS 4.6' section Corrected the URL for GEMM performance measurement documentation. --------- Co-authored-by: Haicheng Wu <57973641+hwu36@users.noreply.github.com> --- CHANGELOG.md | 14 +- README.md | 13 +- include/cutlass/version.h | 2 +- python/CuTeDSL/_mlir_helpers/op.py | 73 +++--- .../CuTeDSL/cutlass/base_dsl/env_manager.py | 2 + .../CuTeDSL/cutlass/base_dsl/runtime/cuda.py | 24 +- python/CuTeDSL/cutlass/cute/__init__.py | 4 + .../cutlass/cute/arch/nvvm_wrappers.py | 22 +- python/CuTeDSL/cutlass/cute/core.py | 152 +++++++++++- python/CuTeDSL/cutlass/jax/__init__.py | 6 + python/CuTeDSL/cutlass/jax/ffi.py | 223 ++++++++++++------ python/CuTeDSL/cutlass/jax/primitive.py | 30 ++- python/CuTeDSL/cutlass/testing.py | 22 +- python/CuTeDSL/requirements-cu13.txt | 2 +- python/CuTeDSL/requirements.txt | 2 +- python/cutlass_cppgen/__init__.py | 2 +- python/setup_cutlass.py | 2 +- python/setup_library.py | 2 +- python/setup_pycute.py | 2 +- 19 files changed, 458 insertions(+), 141 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d0b798b9b..d0099c58d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,9 +2,21 @@ # CUTLASS 4.x +## [4.6.1](https://github.com/NVIDIA/cutlass/releases/tag/v4.6.1) (2026-07-13) + +* Bug fixing and improvements + - Fixed following issues: + - https://github.com/NVIDIA/cutlass/issues/3243 + - https://github.com/NVIDIA/cutlass/issues/3359 + - https://github.com/NVIDIA/cutlass/issues/3365 + - https://github.com/NVIDIA/cutlass/issues/3312 + - Fixed a compilation failure on Thor with 12.9 wheel + - Fixed a per regression of flash attention v2 on Ampere + - Allow custom and multiple FFI call registration for Jax + ## [4.6.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.6.0) (2026-07-01) -* Release [documentation](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/gemm_performance_measurement_methodology_guidelines.md) that explains how to accurately profiling GEMM performance. +* Release [documentation](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/gemm_performance_measurement_methodology_guidelines.html) that explains how to accurately profiling GEMM performance. ### CuTe DSL * New features diff --git a/README.md b/README.md index 6ea198c29..ba7ea83e4 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,9 @@ ![ALT](./media/images/gemm-hierarchy-with-epilogue-no-labels.png "Complete CUDA GEMM decomposition") # Overview -# CUTLASS 4.6.0 +# CUTLASS 4.6.1 -_CUTLASS 4.6.0 - July 2026_ +_CUTLASS 4.6.1 - July 2026_ CUTLASS is a collection of abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and related computations at all levels and scales within CUDA. It incorporates strategies for @@ -45,7 +45,7 @@ To get started quickly - please refer : # What's New in CUTLASS 4.6 -* Release [documentation](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/gemm_performance_measurement_methodology_guidelines.md) that explains how to accurately profiling GEMM performance. +* Release [documentation](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/gemm_performance_measurement_methodology_guidelines.html) that explains how to accurately profiling GEMM performance. ## CuTe DSL * New features @@ -78,6 +78,9 @@ To get started quickly - please refer : - cute.core.ThrMma, please use cute.ThrMma instead - cute.core.ThrCopy, please use cute.ThrCopy instead - cute.make_fragment, please use cute.make_rmem_tensor instead + - Fixed a compilation failure on Thor with 12.9 wheel + - Fixed a per regression of flash attention v2 on Ampere + - Allow custom and multiple FFI call registration for Jax - Fixed following issues - https://github.com/NVIDIA/cutlass/issues/3132 - https://github.com/NVIDIA/cutlass/issues/3170 @@ -89,6 +92,10 @@ To get started quickly - please refer : - https://github.com/NVIDIA/cutlass/issues/3329 - https://github.com/NVIDIA/cutlass/issues/3342 - https://github.com/NVIDIA/cutlass/issues/3348 + - https://github.com/NVIDIA/cutlass/issues/3243 + - https://github.com/NVIDIA/cutlass/issues/3359 + - https://github.com/NVIDIA/cutlass/issues/3365 + - https://github.com/NVIDIA/cutlass/issues/3312 ### CUTLASS Operator API * CUTLASS Operator API is a new addition to the CUTLASS Python stack, providing easy interfaces diff --git a/include/cutlass/version.h b/include/cutlass/version.h index 615a53104..a0c27d9fe 100644 --- a/include/cutlass/version.h +++ b/include/cutlass/version.h @@ -40,7 +40,7 @@ #define CUTLASS_MAJOR 4 #define CUTLASS_MINOR 6 -#define CUTLASS_PATCH 0 +#define CUTLASS_PATCH 1 #ifdef CUTLASS_VERSIONS_GENERATED #include "cutlass/version_extended.h" diff --git a/python/CuTeDSL/_mlir_helpers/op.py b/python/CuTeDSL/_mlir_helpers/op.py index 8ed06b069..632eb38d3 100644 --- a/python/CuTeDSL/_mlir_helpers/op.py +++ b/python/CuTeDSL/_mlir_helpers/op.py @@ -192,12 +192,49 @@ def _find_user_frame(start_frame: types.FrameType | None) -> types.FrameType | N return start_frame +def _get_caller_frame_info() -> inspect.Traceback | None: + cur_frame = inspect.currentframe() + if cur_frame is None: + return None + wrapper_frame = cur_frame.f_back + start_frame = wrapper_frame.f_back if wrapper_frame is not None else None + frame = _find_user_frame(start_frame) + del cur_frame + if frame is None: + return None + return inspect.getframeinfo(frame) + + +def _get_location_from_frame_info(frameInfo: inspect.Traceback) -> ir.Location: + # In Python < 3.11, getframeinfo returns a NamedTuple without positions. + if not hasattr(frameInfo, "positions"): + file_loc = ir.Location.file( + frameInfo.filename, + frameInfo.lineno, + 0, + ) + else: + file_loc = ir.Location.file( + frameInfo.filename, + frameInfo.positions.lineno, # type: ignore[attr-defined] + frameInfo.positions.col_offset or 0, # type: ignore[attr-defined] + ) + return ir.Location.name( + ( + "".join([c.strip() for c in frameInfo.code_context]) + if frameInfo.code_context + else frameInfo.function + ), + childLoc=file_loc, + ) + + def dsl_user_op(opFunc: Callable[..., Any]) -> Callable[..., Any]: """Decorator for user-facing DSL op wrappers. Responsibilities: - 1. Attach a source-location to every MLIR op built by ``opFunc`` so that + 1. Attach source locations when line info / diagnostics are enabled so diagnostics and IR dumps point back at the user's Python call site. 2. Run trace-time MLIR verification on each newly-built op so verifier errors surface at the call site rather than at module-verify time. @@ -220,40 +257,16 @@ def dsl_user_op(opFunc: Callable[..., Any]) -> Callable[..., Any]: @wraps(opFunc) def wrapper(*args: Any, **kwargs: Any) -> Any: # Pop loc= from kwargs so callers that still pass it don't break. - # We no longer forward it — LOC_TRACEBACKS captures full stacks automatically. + # The wrapper replaces it only when source-location tracking is enabled. loc: Any = kwargs.pop("loc", None) frameInfo = None verifier_error = False - if loc is None and ir.Context.current is not None: - cur_frame = inspect.currentframe() - assert cur_frame is not None - frame = _find_user_frame(cur_frame.f_back) - del cur_frame # break self-ref - assert frame is not None - frameInfo = inspect.getframeinfo(frame) + if loc is None and ir.Context.current is not None and _ENABLE_FRAME_FILTERING: + frameInfo = _get_caller_frame_info() try: - # In Python < 3.11, getframeinfo returns a NamedTuple without positions - if not hasattr(frameInfo, "positions"): - file_loc = ir.Location.file( - frameInfo.filename, - frameInfo.lineno, - 0, - ) - else: - file_loc = ir.Location.file( - frameInfo.filename, - frameInfo.positions.lineno, # type: ignore[attr-defined] - frameInfo.positions.col_offset or 0, # type: ignore[attr-defined] - ) - loc = ir.Location.name( - ( - "".join([c.strip() for c in frameInfo.code_context]) - if frameInfo.code_context - else frameInfo.function - ), - childLoc=file_loc, - ) + if frameInfo is not None: + loc = _get_location_from_frame_info(frameInfo) except RuntimeError: # No MLIR context available (e.g. validation-only call # outside a kernel). Proceed with loc=None so that the diff --git a/python/CuTeDSL/cutlass/base_dsl/env_manager.py b/python/CuTeDSL/cutlass/base_dsl/env_manager.py index e54a4a686..4cb86a08a 100644 --- a/python/CuTeDSL/cutlass/base_dsl/env_manager.py +++ b/python/CuTeDSL/cutlass/base_dsl/env_manager.py @@ -47,6 +47,7 @@ _KEEP_ALL_TOKENS: frozenset[str] = frozenset( "ir-debug", "ptx", "cubin", + "sass", } ) # "all" is a convenience alias that expands to every token above. @@ -78,6 +79,7 @@ def _parse_keep_tokens(raw: str, prefix: str = "") -> frozenset[str]: ir-debug — Raw IR before any passes (old KEEP_IR=1 behaviour) ptx — PTX assembly cubin — CUBIN binary + sass — SASS disassembly """ tokens = frozenset(t.strip().lower() for t in raw.split(",") if t.strip()) if "all" in tokens: diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py index 91f72381b..baba20bf5 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py @@ -228,18 +228,22 @@ class DeviceInfo: def pretty_str(self) -> str: """ Convert DeviceInfo to a formatted string for display. + :return: The formatted string. :rtype: str - Example: - On success: - CUDA devices available: (current: ) - - Architecture: () - - Compatible SM archs: - - Total Memory: GB - On failure: - 1. CUDA initialization failed - 2. Failed to get GPU info: - 3. No devices available + + Example success output:: + + - CUDA devices available: (current: ) + - Architecture: () + - Compatible SM archs: + - Total Memory: GB + + Example failure output:: + + - CUDA initialization failed + - Failed to get GPU info: + - No devices available """ info = "" diff --git a/python/CuTeDSL/cutlass/cute/__init__.py b/python/CuTeDSL/cutlass/cute/__init__.py index 89777742c..d03f6f112 100644 --- a/python/CuTeDSL/cutlass/cute/__init__.py +++ b/python/CuTeDSL/cutlass/cute/__init__.py @@ -132,7 +132,9 @@ from .core import ( Ratio, # FastDivmod operations FastDivmodDivisor, + FastDivmodDivisorV2, fast_divmod_create_divisor, + fast_divmod_create_divisor_v2, basis_value, basis_get, nullspace, @@ -422,7 +424,9 @@ __all__ = [ "union", # FastDivmod operations "FastDivmodDivisor", + "FastDivmodDivisorV2", "fast_divmod_create_divisor", + "fast_divmod_create_divisor_v2", # Modules "arch", "export", diff --git a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py index a9f0573a2..25cd21560 100644 --- a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py +++ b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py @@ -1422,6 +1422,8 @@ def fmax( b: Union[float, Float32], *, abs: bool = False, + nan: bool = False, + ftz: bool = False, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> Float32: @@ -1429,13 +1431,20 @@ def fmax( :param abs: When True, lower to the xorsign-abs form ``sign(a ^ b) * max(|a|, |b|)`` (FMNMX.XORSIGN.ABS); default False is a - plain max. NaN / signed-zero behavior follows the underlying NVVM op. + plain max. + :param nan: When True, propagate NaN following IEEE 754 ``maximum`` + (FMNMX.NaN); default False keeps the NaN-quiet ``maximumNumber`` + behavior of the underlying NVVM op. + :param ftz: When True, flush denormal inputs and outputs to zero + (FMNMX.FTZ); default False preserves denormals. """ return Float32( nvvm.fmax( Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), abs=abs, + nan=nan, + ftz=ftz, loc=loc, ip=ip, ) @@ -1448,6 +1457,8 @@ def fmin( b: Union[float, Float32], *, abs: bool = False, + nan: bool = False, + ftz: bool = False, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None, ) -> Float32: @@ -1455,13 +1466,20 @@ def fmin( :param abs: When True, lower to the xorsign-abs form ``sign(a ^ b) * min(|a|, |b|)`` (FMNMX.XORSIGN.ABS); default False is a - plain min. NaN / signed-zero behavior follows the underlying NVVM op. + plain min. + :param nan: When True, propagate NaN following IEEE 754 ``minimum`` + (FMNMX.NaN); default False keeps the NaN-quiet ``minimumNumber`` + behavior of the underlying NVVM op. + :param ftz: When True, flush denormal inputs and outputs to zero + (FMNMX.FTZ); default False preserves denormals. """ return Float32( nvvm.fmin( Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), abs=abs, + nan=nan, + ftz=ftz, loc=loc, ip=ip, ) diff --git a/python/CuTeDSL/cutlass/cute/core.py b/python/CuTeDSL/cutlass/cute/core.py index a3a765148..a81af1bc0 100644 --- a/python/CuTeDSL/cutlass/cute/core.py +++ b/python/CuTeDSL/cutlass/cute/core.py @@ -14,6 +14,7 @@ from collections.abc import Iterable from functools import partial, reduce import inspect from inspect import isclass +import warnings from typing import ( Any, Callable, @@ -442,9 +443,15 @@ class IntValue(cutlass_arith.ArithValue): ip: Optional[ir.InsertionPoint] = None, ) -> ir.Value: if isinstance(self.type, ir.IntegerType): - def_op = self.owner.operation - if def_op.name == "cute.get_scalars": - return def_op.operands[0] + # A block argument (e.g. a kernel parameter, such as the scalar + # divisor a FastDivmodDivisorV2 carries across the kernel boundary) + # is owned by an ir.Block, not an ir.Operation, so it has no + # defining op and no cute.get_scalars shortcut. Fall through to + # wrap it as an int_tuple in that case. + if not isinstance(self.owner, ir.Block): + def_op = self.owner.operation + if def_op.name == "cute.get_scalars": + return def_op.operands[0] assert not isinstance(self.type, _cute_ir.IntTupleType) @@ -2174,7 +2181,17 @@ def printf( raise TypeError(f"unsupported argument type in printf, got {type(arg)}") processed_args = [process_arg(a) for a in args] - _cute_ir.print_(processed_args, fmt=fmt, loc=loc, ip=ip) + operand_signed = [bool(getattr(type(a), "signed", True)) for a in args] + if all(operand_signed): + _cute_ir.print_(processed_args, fmt=fmt, loc=loc, ip=ip) + else: + _cute_ir.print_( + processed_args, + fmt=fmt, + operand_signed=ir.DenseBoolArrayAttr.get(operand_signed), + loc=loc, + ip=ip, + ) @dsl_user_op @@ -6262,6 +6279,13 @@ class FastDivmodDivisor: This class wraps a FastDivmod divisor and enables natural Python operator syntax. + .. deprecated:: + Use :class:`FastDivmodDivisorV2` instead. V2 additionally carries the + scalar divisor across kernel boundaries (2 MLIR values per object + instead of 1), so ``.divisor`` is readable inside kernels; + arithmetic is unchanged. This class keeps the legacy 1-value + serialization contract for existing integrations. + :ivar divisor: The original divisor value (publicly accessible) :ivar _divisor_mlir: The FastDivmod divisor MLIR value (internal) @@ -6290,6 +6314,20 @@ class FastDivmodDivisor: :param is_power_of_2: Whether divisor is known to be a power of 2. Defaults to False. """ + # Subclasses (FastDivmodDivisorV2) share this __init__; only direct + # use of the legacy class is deprecated. + if type(self) is FastDivmodDivisor: + warnings.warn( + "FastDivmodDivisor is deprecated in favor of " + "cute.FastDivmodDivisorV2 / cute.fast_divmod_create_divisor_v2. " + "V2 additionally carries the scalar divisor across kernel " + "boundaries (2 MLIR values per object instead of 1), so " + "'.divisor' is readable inside kernels; " + "arithmetic (divmod, //, %) is unchanged.", + DeprecationWarning, + stacklevel=2, + ) + # Store the original divisor value for public access self._original_divisor = divisor @@ -6407,6 +6445,13 @@ class FastDivmodDivisor: batch_fdd = cute.fast_divmod_create_divisor(batch_size) print(f"Divisor: {batch_fdd.divisor}") # Access the divisor value some_function(divisor=batch_fdd.divisor) # Pass to other functions + + .. note:: + After this object crosses a kernel boundary (e.g. stored in a + params structure passed to a ``@cute.kernel``), the returned value + still references host-side SSA and fails MLIR region isolation if + used inside the kernel (OSS issue #3243). Use + :class:`FastDivmodDivisorV2` to read the divisor inside a kernel. """ return self._original_divisor @@ -6454,7 +6499,7 @@ class FastDivmodDivisor: return new_obj def __repr__(self) -> str: - return f"FastDivmodDivisor(divisor={self._original_divisor}, type={self._divisor_mlir.type})" + return f"{type(self).__name__}(divisor={self._original_divisor}, type={self._divisor_mlir.type})" # Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator @@ -6503,3 +6548,100 @@ def fast_divmod_create_divisor( remainder = linear_idx % divisor """ return FastDivmodDivisor(divisor, loc=loc, ip=ip) + + +class FastDivmodDivisorV2(FastDivmodDivisor): + """ + FastDivmod divisor whose ``.divisor`` property is readable inside kernels. + + Same arithmetic behavior as :class:`FastDivmodDivisor` (``divmod``, ``//``, + ``%``), but serializes **two** MLIR values across region boundaries — the + encoded FastDivmod plus the scalar divisor — so ``.divisor`` resolves to + in-region SSA after the object crosses a kernel boundary (OSS issue #3243): + + .. code-block:: python + + @dataclass + class Params: + fdd: cute.FastDivmodDivisorV2 + + @cute.kernel + def kernel(out: cute.Tensor, params: Params): + out[0] = params.fdd.divisor # OK: region-local SSA + + :class:`FastDivmodDivisor` keeps the legacy 1-value serialization contract + for backward compatibility; its ``.divisor`` is not readable inside a + kernel. + """ + + def __extract_mlir_values__(self) -> List[ir.Value]: + """Extract MLIR values for Host->Device transfer. + + Two SSA values are emitted: the encoded FastDivmod (``_divisor_mlir``) + and the scalar divisor that was used to build it. The encoded value + still flows through GridInvariantCodeMotionPass for host hoisting; the + scalar value is needed so ``.divisor`` resolves to in-region SSA after + crossing the kernel boundary (issue #3243). + """ + divisor_for_pack = self._original_divisor + if isinstance(divisor_for_pack, ir.Value): + divisor_ir = divisor_for_pack + else: + divisor_ir = Int32(divisor_for_pack).ir_value() + return [self._divisor_mlir, divisor_ir] + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "FastDivmodDivisorV2": + """Reconstruct FastDivmodDivisorV2 from MLIR values. + + Rebuilds ``_original_divisor`` from the SSA passed in ``values[1]`` so + that ``.divisor`` reads kernel-region SSA, not the host-side template. + """ + if len(values) != 2: + raise ValueError( + "FastDivmodDivisorV2 expects exactly 2 MLIR values (encoded " + f"divisor + scalar divisor SSA), got {len(values)}. If this " + "object is held by a params class with a hand-written " + "__new_from_mlir_values__, make sure it slices 2 values per " + "FastDivmodDivisorV2 field." + ) + new_obj = object.__new__(FastDivmodDivisorV2) + new_obj._divisor_mlir = values[0] + # values[1] may arrive as a raw ir.Value or as a typed integer wrapper + # (the framework value caster reconstructs typed wrappers across the + # kernel boundary). Normalize to IntValue so '.divisor' supports + # arithmetic and repr inside the kernel. + scalar_divisor = values[1] + if isinstance(scalar_divisor, ir.Value): + scalar_divisor = IntValue(scalar_divisor) + new_obj._original_divisor = scalar_divisor + return new_obj + + +@dsl_user_op +def fast_divmod_create_divisor_v2( + divisor: Integer, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> FastDivmodDivisorV2: + """Create a FastDivmod divisor whose ``.divisor`` is readable inside kernels. + + Behaves like :func:`fast_divmod_create_divisor`, but the returned + :class:`FastDivmodDivisorV2` serializes both the encoded FastDivmod and the + scalar divisor across kernel boundaries, so ``.divisor`` resolves to + region-local SSA inside a kernel (OSS issue #3243). + + :param divisor: The divisor value (should be runtime-dynamic value) + :type divisor: Integer + :return: FastDivmodDivisorV2 object with operator overloading support + :rtype: FastDivmodDivisorV2 + + **Example:** + + .. code-block:: python + + divisor = fast_divmod_create_divisor_v2(batch_size) + quotient, remainder = divmod(linear_idx, divisor) + d = divisor.divisor # readable on host AND inside kernels + """ + return FastDivmodDivisorV2(divisor, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/jax/__init__.py b/python/CuTeDSL/cutlass/jax/__init__.py index 42a0e49c6..30d56d8ca 100644 --- a/python/CuTeDSL/cutlass/jax/__init__.py +++ b/python/CuTeDSL/cutlass/jax/__init__.py @@ -65,7 +65,10 @@ if is_available(): get_export_disabled_safety_checks, find_cute_dsl_runtime_library, register_ffi, + set_ffi_call_targets, + disable_automatic_ffi_registration, is_ffi_registered, + get_cutlass_call_ffi_name, get_cutlass_call_ffi_version, ) from . import testing @@ -88,6 +91,9 @@ if is_available(): "get_export_disabled_safety_checks", "is_ffi_registered", "register_ffi", + "set_ffi_call_targets", + "disable_automatic_ffi_registration", + "get_cutlass_call_ffi_name", "get_cutlass_call_ffi_version", "is_available", "testing", diff --git a/python/CuTeDSL/cutlass/jax/ffi.py b/python/CuTeDSL/cutlass/jax/ffi.py index 7d27f2a2f..6d6eb419e 100644 --- a/python/CuTeDSL/cutlass/jax/ffi.py +++ b/python/CuTeDSL/cutlass/jax/ffi.py @@ -9,9 +9,10 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import Any, Optional, Sequence +from typing import Any, Sequence from pathlib import Path from functools import cache +from threading import Lock import logging import ctypes @@ -26,66 +27,72 @@ logger = logging.getLogger(__name__) _CUTE_DSL_RUNTIME_LIBRARY_NAME = "cute_dsl_runtime" -# V1 targets for older jax clients -_CUTLASS_CALL_TARGETS_V1 = { - "CuteDSLRT_NvJaxCutlassCall": { - "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v1", - "execute": "CuteDSLRT_NvJaxCutlassCallExecute_v1", +_JAX_FFI_V2_MIN_VERSION = (0, 9, 1) + +_CUTLASS_CALL_TARGETS = { + # V1 targets for older JAX clients. + 1: { + "CuteDSLRT_NvJaxCutlassCall_v1": { + "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v1", + "execute": "CuteDSLRT_NvJaxCutlassCallExecute_v1", + }, + "CuteDSLRT_NvJaxCutlassCallNoCudaGraph_v1": { + "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v1", + "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph_v1", + }, }, - "CuteDSLRT_NvJaxCutlassCallNoCudaGraph": { - "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v1", - "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph_v1", + # V2 targets for newer JAX clients supporting stateful FFI calls. + 2: { + "CuteDSLRT_NvJaxCutlassCall_v2": { + "execute": "CuteDSLRT_NvJaxCutlassCallExecute_v2", + "instantiate": "CuteDSLRT_NvJaxCutlassCallInstantiate_v2", + "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v2", + }, + "CuteDSLRT_NvJaxCutlassCallNoCudaGraph_v2": { + "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph_v2", + "instantiate": "CuteDSLRT_NvJaxCutlassCallInstantiate_v2", + "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v2", + }, + }, +} +_CUTLASS_CALL_TYPES = { + 1: {}, + 2: { + "CuteDSLRT_NvJaxCutlassCallTypes_v2": { + "type_id": "CuteDSLRT_NvJaxCutlassCallStateTypeId_v2", + "type_info": "CuteDSLRT_NvJaxCutlassCallStateTypeInfo_v2", + } }, } -# V2 targets for newer jax clients supporting stateful FFI calls. -_JAX_FFI_V2_MIN_VERSION = (0, 9, 1) -_CUTLASS_CALL_TARGETS_V2 = { - "CuteDSLRT_NvJaxCutlassCall": { - "execute": "CuteDSLRT_NvJaxCutlassCallExecute_v2", - "instantiate": "CuteDSLRT_NvJaxCutlassCallInstantiate_v2", - "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v2", - }, - "CuteDSLRT_NvJaxCutlassCallNoCudaGraph": { - "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph_v2", - "instantiate": "CuteDSLRT_NvJaxCutlassCallInstantiate_v2", - "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v2", - }, -} -_CUTLASS_CALL_TYPES_V2 = { - "CuteDSLRT_NvJaxCutlassCallTypes": { - "type_id": "CuteDSLRT_NvJaxCutlassCallStateTypeId_v2", - "type_info": "CuteDSLRT_NvJaxCutlassCallStateTypeInfo_v2", - } -} +_FFI_REGISTRATION_LOCK = Lock() +_REGISTERED_FFI_CALL_TARGETS: set[str] = set() +_REGISTERED_FFI_TYPES: set[str] = set() +_AUTOMATIC_FFI_REGISTRATION_ENABLED = True +_FFI_CALL_TARGETS_OVERRIDE: tuple[str, str] | None = None def get_cutlass_call_ffi_version() -> int: - """Returns the FFI API version based on JAX version.""" - if jax.version.__version_info__ >= _JAX_FFI_V2_MIN_VERSION: - return 2 - else: - return 1 + """Returns the FFI version supported by the installed JAX.""" + return 2 if jax.version.__version_info__ >= _JAX_FFI_V2_MIN_VERSION else 1 -def get_cutlass_call_ffi_name(allow_cuda_graph: bool) -> str: - """Returns the FFI target to call when running cutlass_call functions.""" - if allow_cuda_graph: - return "CuteDSLRT_NvJaxCutlassCall" - else: - return "CuteDSLRT_NvJaxCutlassCallNoCudaGraph" +def get_cutlass_call_ffi_name(allow_cuda_graph: bool, version: int) -> str: + """Returns the FFI target for ``version``.""" + suffix = "" if allow_cuda_graph else "NoCudaGraph" + return f"CuteDSLRT_NvJaxCutlassCall{suffix}_v{version}" def get_export_disabled_safety_checks() -> Sequence[jax.export.DisabledSafetyCheck]: - """Returns jax.export.DisabledSafetyCheck to allow cutlass_call kernels.""" - targets = set(_CUTLASS_CALL_TARGETS_V1.keys()) | set( - _CUTLASS_CALL_TARGETS_V2.keys() + """Returns export safety checks for the built-in FFI targets.""" + targets = set().union(*_CUTLASS_CALL_TARGETS.values()) + return tuple( + jax.export.DisabledSafetyCheck.custom_call(target) for target in sorted(targets) ) - return tuple([jax.export.DisabledSafetyCheck.custom_call(t) for t in targets]) @cache -def find_cute_dsl_runtime_library() -> Optional[str]: +def find_cute_dsl_runtime_library() -> str | None: """Searches for the CuTeDSL runtime library.""" dsl = CuTeDSL._get_dsl() candidate_libs = [] @@ -125,17 +132,25 @@ def find_cute_dsl_runtime_library() -> Optional[str]: return None -_FFI_CALLS_REGISTERED = False +def _is_ffi_version_registered(version: int) -> bool: + return all( + name in _REGISTERED_FFI_CALL_TARGETS for name in _CUTLASS_CALL_TARGETS[version] + ) and all(name in _REGISTERED_FFI_TYPES for name in _CUTLASS_CALL_TYPES[version]) -def register_ffi(ffi_version: int = get_cutlass_call_ffi_version()) -> None: - """Registers custom calls with Jax/XLA runtime. - - A specific version can be requested using `ffi_version` argument. Attempting - to register non default FFI versions may not work with your specific JAX. - """ - global _FFI_CALLS_REGISTERED - if _FFI_CALLS_REGISTERED: +def _register_ffi(version: int) -> None: + """Registers ``version`` while ``_FFI_REGISTRATION_LOCK`` is held.""" + call_targets = { + name: symbols + for name, symbols in _CUTLASS_CALL_TARGETS[version].items() + if name not in _REGISTERED_FFI_CALL_TARGETS + } + type_targets = { + name: symbols + for name, symbols in _CUTLASS_CALL_TYPES[version].items() + if name not in _REGISTERED_FFI_TYPES + } + if not call_targets and not type_targets: return runtime_library = find_cute_dsl_runtime_library() @@ -158,32 +173,96 @@ def register_ffi(ffi_version: int = get_cutlass_call_ffi_version()) -> None: handler[stage] = jax.ffi.pycapsule(fn) logger.debug(f"Registering ffi handler: {target_name}, {handler}") jax.ffi.register_ffi_target(target_name, handler, platform="CUDA") + _REGISTERED_FFI_CALL_TARGETS.add(target_name) def _register_ffi_types(lib: ctypes.CDLL, types: dict[str, dict[str, str]]) -> None: - for type_name, type_dict_targets in types.items(): + for type_name, type_targets in types.items(): type_dict: dict[str, Any] = {} - for field, fn_name in type_dict_targets.items(): + for field, fn_name in type_targets.items(): fn = getattr(lib, fn_name) fn.restype = ctypes.c_void_p type_dict[field] = jax.ffi.pycapsule(fn()) logger.debug(f"Registering ffi type: {type_name}, {type_dict}") jax.ffi.register_ffi_type(type_name, type_dict, platform="CUDA") # type: ignore[arg-type] + _REGISTERED_FFI_TYPES.add(type_name) - # Register the custom FFI targets. - match ffi_version: - case 1: - _register_ffi_targets(lib, _CUTLASS_CALL_TARGETS_V1) - # no types for v1 - case 2: - _register_ffi_types(lib, _CUTLASS_CALL_TYPES_V2) - _register_ffi_targets(lib, _CUTLASS_CALL_TARGETS_V2) - case _: - raise ValueError(f"Invalid FFI version {ffi_version}") - - _FFI_CALLS_REGISTERED = True + _register_ffi_types(lib, type_targets) + _register_ffi_targets(lib, call_targets) -def is_ffi_registered() -> bool: - """Returns true if the FFI calls have been registered with Jax/XLA.""" - global _FFI_CALLS_REGISTERED - return _FFI_CALLS_REGISTERED +def register_ffi(version: int | None = None) -> None: + """Explicitly registers a built-in FFI version. + + If ``version`` is omitted, the version supported by the installed JAX is + registered. Multiple supported versions may be registered in one process. + """ + if version is None: + version = get_cutlass_call_ffi_version() + + with _FFI_REGISTRATION_LOCK: + _register_ffi(version) + + +def set_ffi_call_targets( + default_call_target: str, + no_cuda_graph_call_target: str | None = None, +) -> None: + """Overrides the process-wide targets used by future ``cutlass_call`` objects. + + If ``no_cuda_graph_call_target`` is omitted, ``default_call_target`` is used + for both CUDA graph modes. CuTeDSL does not register these targets or add + export safety checks for them; the caller is responsible for both. Explicit + :func:`register_ffi` calls continue to register built-in targets. + """ + if no_cuda_graph_call_target is None: + no_cuda_graph_call_target = default_call_target + + global _FFI_CALL_TARGETS_OVERRIDE + with _FFI_REGISTRATION_LOCK: + _FFI_CALL_TARGETS_OVERRIDE = ( + default_call_target, + no_cuda_graph_call_target, + ) + + +def _register_and_get_default_ffi_call_target(allow_cuda_graph: bool) -> str: + """Returns and, unless disabled, registers the JAX-compatible FFI target.""" + with _FFI_REGISTRATION_LOCK: + if _FFI_CALL_TARGETS_OVERRIDE is not None: + return _FFI_CALL_TARGETS_OVERRIDE[0 if allow_cuda_graph else 1] + + version = get_cutlass_call_ffi_version() + if _AUTOMATIC_FFI_REGISTRATION_ENABLED: + _register_ffi(version) + return get_cutlass_call_ffi_name(allow_cuda_graph, version) + + +def disable_automatic_ffi_registration() -> None: + """Disables registration initiated by future ``cutlass_call`` objects. + + Call this before constructing a ``cutlass_call`` to prevent its automatic + registration. + Existing registrations are unchanged, and explicit :func:`register_ffi` + calls remain enabled. This operation is process-wide and idempotent. + """ + global _AUTOMATIC_FFI_REGISTRATION_ENABLED + with _FFI_REGISTRATION_LOCK: + _AUTOMATIC_FFI_REGISTRATION_ENABLED = False + + +def is_ffi_registered(version: int | None = None, *, name: str | None = None) -> bool: + """Returns whether this module registered an FFI version or call/type name. + + ``name`` checks an individual built-in call target or type registered by + this module. Without a selector, this checks the built-in version supported + by the installed JAX. + """ + if version is not None and name is not None: + raise ValueError("'version' and 'name' cannot both be specified") + if name is not None: + with _FFI_REGISTRATION_LOCK: + return name in _REGISTERED_FFI_CALL_TARGETS or name in _REGISTERED_FFI_TYPES + if version is None: + version = get_cutlass_call_ffi_version() + with _FFI_REGISTRATION_LOCK: + return _is_ffi_version_registered(version) diff --git a/python/CuTeDSL/cutlass/jax/primitive.py b/python/CuTeDSL/cutlass/jax/primitive.py index ff37b1c84..ae94e1db7 100644 --- a/python/CuTeDSL/cutlass/jax/primitive.py +++ b/python/CuTeDSL/cutlass/jax/primitive.py @@ -26,7 +26,9 @@ from .types import ( default_tensor_spec, TensorSpec, ) -from .ffi import get_cutlass_call_ffi_name, is_ffi_registered, register_ffi +from .ffi import ( + _register_and_get_default_ffi_call_target, +) logger = logging.getLogger(__name__) @@ -45,6 +47,7 @@ def cutlass_call( output_mode: Any = None, input_output_aliases: dict[int, int] | None = None, allow_cuda_graph: bool = True, + ffi_call_target: str | None = None, compile_options: str | None = None, use_static_tensors: bool = False, **kwargs: Any, @@ -91,6 +94,14 @@ def cutlass_call( Indices are into the flattened input and output pytrees. allow_cuda_graph: If ``False``, prevents XLA from capturing this call in a CUDA graph. Defaults to ``True``. + ffi_call_target: Exact FFI target name to call without automatic + registration or default-target selection. + The target must implement the CuTeDSL call ABI, and the caller is + responsible for registering it with JAX. *allow_cuda_graph* does + not modify an explicit target name. When exporting, the caller + must also allow the custom target with + :meth:`jax.export.DisabledSafetyCheck.custom_call`. Defaults to + ``None``. compile_options: Optional string of compiler flags forwarded to ``cute.compile``. use_static_tensors: If ``True``, tensor shapes and strides are baked in @@ -108,6 +119,10 @@ def cutlass_call( """ if output_shape_dtype is None: raise ValueError("'output_shape_dtype' must be specified.") + if ffi_call_target is None: + # Resolve the process default before binding so the target participates + # in JAX's compilation cache key. + ffi_call_target = _register_and_get_default_ffi_call_target(allow_cuda_graph) output_shape_dtype = jax.tree.map( lambda leaf: jax.ShapeDtypeStruct(leaf.shape, leaf.dtype), output_shape_dtype @@ -137,6 +152,7 @@ def cutlass_call( output_spec=output_spec, input_output_aliases=input_output_aliases, allow_cuda_graph=allow_cuda_graph, + ffi_call_target=ffi_call_target, compile_options=compile_options, use_static_tensors=use_static_tensors, **kwargs, @@ -231,6 +247,7 @@ def _cutlass_call_impl( output_spec: Any, input_output_aliases: dict[int, int], allow_cuda_graph: bool, + ffi_call_target: str, compile_options: str | None, use_static_tensors: bool, **kwargs: Any, @@ -261,6 +278,7 @@ def _cutlass_call_impl( output_spec_flat=output_spec_flat, input_output_aliases=tuple(input_output_aliases.items()), allow_cuda_graph=allow_cuda_graph, + ffi_call_target=ffi_call_target, compile_options=compile_options, use_static_tensors=use_static_tensors, **kwargs, @@ -289,6 +307,7 @@ def cutlass_call_inner_p_impl( output_spec_flat: tuple[TensorSpec, ...], input_output_aliases: tuple[tuple[int, int], ...], allow_cuda_graph: bool, + ffi_call_target: str, compile_options: str | None, use_static_tensors: bool, **kwargs: Any, @@ -309,13 +328,6 @@ def cutlass_call_inner_p_impl( kernel = get_or_compile_kernel(fn, spec) - # Ensure our FFI target is registered. We do this lazily here - # so that we only load the dependant library if needed. - if not is_ffi_registered(): - register_ffi() - - call_name = get_cutlass_call_ffi_name(allow_cuda_graph) - # Convert explicit layout constraints from CuTeDSL to JAX order. ``None`` is # passed through intentionally: jax.ffi.ffi_call treats it as default # row-major layout. @@ -323,7 +335,7 @@ def cutlass_call_inner_p_impl( output_layouts = [cutlass_to_jax_layout_order(s.layout) for s in output_spec_flat] fun = jax.ffi.ffi_call( - call_name, + ffi_call_target, result_shape_dtypes=output_shape_dtype_flat, input_output_aliases=dict(spec.input_output_aliases), input_layouts=input_layouts, diff --git a/python/CuTeDSL/cutlass/testing.py b/python/CuTeDSL/cutlass/testing.py index 0a2bd7f75..27c101e07 100644 --- a/python/CuTeDSL/cutlass/testing.py +++ b/python/CuTeDSL/cutlass/testing.py @@ -290,6 +290,22 @@ def _does_kernel_use_stream( return False +def _does_kernel_use_stream_with_jit_retry( + kernel: Callable[..., Any], + stream: cuda_driver.CUstream, + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> bool: + uses_stream = _does_kernel_use_stream(kernel, stream, args, kwargs) + if uses_stream or not hasattr(kernel, "_dsl_cls"): + return uses_stream + + # The first invocation of an uncompiled @cute.jit function can spend the + # capture attempt compiling instead of recording a launch. Retry once with + # the now-compiled callable before reporting a stream mismatch. + return _does_kernel_use_stream(kernel, stream, args, kwargs) + + def benchmark( callable: Callable, *, @@ -401,7 +417,7 @@ def benchmark( if ( not use_cuda_graphs and int(stream) != int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) - and not _does_kernel_use_stream( + and not _does_kernel_use_stream_with_jit_retry( callable, stream, workspaces[0].args, workspaces[0].kwargs ) ): @@ -716,7 +732,9 @@ def _benchmark_for_autotune( if int(current_stream) != int( cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) - ) and not _does_kernel_use_stream(callable, current_stream, args, kwargs): + ) and not _does_kernel_use_stream_with_jit_retry( + callable, current_stream, args, kwargs + ): raise ValueError(f"Incorrect stream passed to kernel: {current_stream}") if use_cold_l2: diff --git a/python/CuTeDSL/requirements-cu13.txt b/python/CuTeDSL/requirements-cu13.txt index 561815af1..21f04a18c 100644 --- a/python/CuTeDSL/requirements-cu13.txt +++ b/python/CuTeDSL/requirements-cu13.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements-cu13.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl[cu13]==4.6.0 +nvidia-cutlass-dsl[cu13]==4.6.1 diff --git a/python/CuTeDSL/requirements.txt b/python/CuTeDSL/requirements.txt index 2f5b51e0f..3e1b736d0 100644 --- a/python/CuTeDSL/requirements.txt +++ b/python/CuTeDSL/requirements.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl==4.6.0 +nvidia-cutlass-dsl==4.6.1 diff --git a/python/cutlass_cppgen/__init__.py b/python/cutlass_cppgen/__init__.py index 807632e52..34ffb8df3 100644 --- a/python/cutlass_cppgen/__init__.py +++ b/python/cutlass_cppgen/__init__.py @@ -133,7 +133,7 @@ def get_option_registry(): this._option_registry = OptionRegistry(device_cc()) return this._option_registry -this.__version__ = '4.6.0' +this.__version__ = '4.6.1' from cutlass_cppgen.backend import create_memory_pool from cutlass_cppgen.emit.pytorch import pytorch diff --git a/python/setup_cutlass.py b/python/setup_cutlass.py index eba123f33..a3051f6bb 100644 --- a/python/setup_cutlass.py +++ b/python/setup_cutlass.py @@ -51,7 +51,7 @@ setup_pycute.perform_setup() setup( name='cutlass_cppgen', - version='4.6.0', + version='4.6.1', description='CUTLASS Pythonic Interface', package_dir={'': '.'}, packages=[ diff --git a/python/setup_library.py b/python/setup_library.py index 353a696bb..8ec1b3edc 100644 --- a/python/setup_library.py +++ b/python/setup_library.py @@ -36,7 +36,7 @@ from setuptools import setup def perform_setup(): setup( name='cutlass_library', - version='4.6.0', + version='4.6.1', description='CUTLASS library generation scripts', packages=['cutlass_library'] ) diff --git a/python/setup_pycute.py b/python/setup_pycute.py index 58bade228..c7f1297c0 100644 --- a/python/setup_pycute.py +++ b/python/setup_pycute.py @@ -36,7 +36,7 @@ from setuptools import setup def perform_setup(): setup( name='pycute', - version='4.6.0', + version='4.6.1', description='Python implementation of CuTe', packages=['pycute'], )