NVIDIA/cutlass
1
0
Fork 0
mirror of https://github.com/NVIDIA/cutlass.git synced 2026-07-17 09:07:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

v4.6.1 update. (#3383)

Browse Source 
* 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>
This commit is contained in:
Junkai-Wu
2026-07-14 10:49:52 +08:00
committed by GitHub
parent 05fd39dca2
commit 1e73948292
19 changed files with 458 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
python/CuTeDSL/_mlir_helpers/op.py
View File

@@ -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

2
python/CuTeDSL/cutlass/base_dsl/env_manager.py
View File

@@ -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:

24
python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py
View File

@@ -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: <device_count> (current: <current_device>)
- Architecture: <arch_name> (<sm_arch>)
- Compatible SM archs: <compatible_archs>
- Total Memory: <memory_gb> GB
On failure:
1. CUDA initialization failed
2. Failed to get GPU info: <error_message>
3. No devices available
Example success output::
- CUDA devices available: <device_count> (current: <current_device>)
- Architecture: <arch_name> (<sm_arch>)
- Compatible SM archs: <compatible_archs>
- Total Memory: <memory_gb> GB
Example failure output::
- CUDA initialization failed
- Failed to get GPU info: <error_message>
- No devices available
"""
info = ""

4
python/CuTeDSL/cutlass/cute/__init__.py
View File

@@ -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",

22
python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py
View File

@@ -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,
)

152
python/CuTeDSL/cutlass/cute/core.py
View File

@@ -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)

6
python/CuTeDSL/cutlass/jax/__init__.py
View File

@@ -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",

223
python/CuTeDSL/cutlass/jax/ffi.py
View File

@@ -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)

30
python/CuTeDSL/cutlass/jax/primitive.py
View File

@@ -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,

22
python/CuTeDSL/cutlass/testing.py
View File

@@ -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:

2
python/CuTeDSL/requirements-cu13.txt
View File

@@ -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

2
python/CuTeDSL/requirements.txt
View File

@@ -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

2
python/cutlass_cppgen/__init__.py
View File

@@ -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

2
python/setup_cutlass.py
View File

@@ -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=[

2
python/setup_library.py
View File

@@ -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']
)

2
python/setup_pycute.py
View File

@@ -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'],
)