leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

v4.3.1 update. (#2817)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-27 22:49:30 +08:00
committed by GitHub
parent 2052fd3885
commit 1de3a576cc
44 changed files with 3316 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
media/docs/pythonDSL/cute_dsl_general/framework_integration.rst
View File

@@ -415,24 +415,36 @@ The following example demonstrates how to use ``mark_compact_shape_dynamic`` to
)
# Layout in DLTensorWrapper has int32 overflow risk. Please set use_32bit_stride to False.
Leveraging TVM FFI for Faster PyTorch Interop
---------------------------------------------
The latest version of |DSL| supports TVM FFI to improve interoperability with PyTorch
and other machine learning frameworks. Using TVM FFI provides the following features:
- Faster JIT function invocation.
- Direct acceptance of ``torch.Tensor`` objects as function arguments.
- Enhanced error handling and kernel validation.
- Seamless integration with multiple programming languages.
For more details, see :doc:`compile_with_tvm_ffi`.
Bypass the DLPack Protocol
--------------------------
In certain scenarios, users may wish to bypass the DLPack protocol and invoke the JIT function directly.
This can be accomplished by creating a lightweight JIT wrapper around the existing JIT function,
In certain scenarios, users may wish to bypass the DLPack protocol and invoke the JIT function directly.
This can be accomplished by creating a lightweight JIT wrapper around the existing JIT function,
utilizing ``cute.ptr`` and ``cute.make_tensor`` to pass pointers and construct tensors directly.
Typical use cases for bypassing DLPack include:
1. Users want to call the JIT function directly to avoid the overhead introduced by the DLPack protocol.
2. DLPack canonicalizes the stride of shape-1 dimensions to 1, which may result in incorrect alignment
2. DLPack canonicalizes the stride of shape-1 dimensions to 1, which may result in incorrect alignment
propagation and affect memory access or performance.
3. DLPack may lack support for some narrow data types.
The following example illustrates how to bypass the DLPack protocol when invoking a JIT function.
Assume we have a pre-defined ``TensorOpGemm`` kernel whose JIT interface expects three
arguments of type ``cute.Tensor``. To enable direct invocation without DLPack, we first define a JIT wrapper
function that accepts ``cute.Pointer`` types as parameters. Within this wrapper, we use ``cute.make_tensor``
Assume we have a pre-defined ``TensorOpGemm`` kernel whose JIT interface expects three
arguments of type ``cute.Tensor``. To enable direct invocation without DLPack, we first define a JIT wrapper
function that accepts ``cute.Pointer`` types as parameters. Within this wrapper, we use ``cute.make_tensor``
to construct tensors from the provided pointers, and then call the ``TensorOpGemm`` kernel as usual.
.. code-block:: python
@@ -459,7 +471,7 @@ to construct tensors from the provided pointers, and then call the ``TensorOpGem
mA = cute.make_tensor(a_ptr, layout=a_layout)
mB = cute.make_tensor(b_ptr, layout=b_layout)
mC = cute.make_tensor(c_ptr, layout=c_layout)
# TensorOpGemm is a pre-defined kernel from our example
tensor_op_gemm = TensorOpGemm(
a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1)
@@ -467,9 +479,9 @@ to construct tensors from the provided pointers, and then call the ``TensorOpGem
tensor_op_gemm(mA, mB, mC)
To pass a PyTorch tensor to this new JIT wrapper, we retrieve the raw pointer from the PyTorch tensor
To pass a PyTorch tensor to this new JIT wrapper, we retrieve the raw pointer from the PyTorch tensor
and create a ``cute.Pointer`` instance using ``cute.make_ptr``.
This approach allows us to bypass the DLPack protocol entirely, avoiding its overhead and potential
This approach allows us to bypass the DLPack protocol entirely, avoiding its overhead and potential
issues with shape-1 dimension handling.
.. code-block:: python
@@ -483,7 +495,7 @@ issues with shape-1 dimension handling.
c = torch.randn(
n, m, l, dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
# from cutlass.cute.runtime import make_ptr
a_ptr = make_ptr(
cutlass.Float16, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
@@ -495,3 +507,5 @@ issues with shape-1 dimension handling.
cutlass.Float16, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
tensor_op_gemm_wrapper(a_ptr, b_ptr, c_ptr, m, n, k, l)