diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
index d5b48f45..979609c3 100644
--- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
@@ -1715,7 +1715,6 @@ def compile_bmm(
max_active_clusters: cutlass.Constexpr,
stream: cuda.CUstream,
epilogue_op: cutlass.Constexpr = lambda x: x,
- options: str = "",
):
from cutlass.cute.runtime import make_fake_compact_tensor
@@ -1749,7 +1748,15 @@ def compile_bmm(
)
return cute.compile(
- bmm, gemm_op, a, b, c, max_active_clusters, stream, epilogue_op, options=options
+ bmm,
+ gemm_op,
+ a,
+ b,
+ c,
+ max_active_clusters,
+ stream,
+ epilogue_op,
+ options="--enable-tvm-ffi",
)
@@ -1811,7 +1818,6 @@ def run(
iterations: int = 1,
skip_ref_check: bool = False,
use_cold_l2: bool = False,
- use_tvm_ffi: bool = False,
benchmark: bool = False,
**kwargs,
):
@@ -1853,8 +1859,6 @@ def run(
:type skip_ref_check: bool, optional
:param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False.
:type use_cold_l2: bool, optional
- :param use_tvm_ffi: Whether to use TVM FFI for the kernel, defaults to False.
- :type use_tvm_ffi: bool, optional
:param benchmark: Whether to only benchmark the kernel, defaults to False.
:type benchmark: bool, optional
:raises RuntimeError: If CUDA GPU is not available.
@@ -1874,7 +1878,7 @@ def run(
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
- print(f"Use TVM FFI: {'True' if use_tvm_ffi else 'False'}")
+ print(f"Use TVM FFI")
import torch
from cutlass.torch import dtype as torch_dtype
@@ -1906,10 +1910,6 @@ def run(
cluster_shape_mn[0] * cluster_shape_mn[1]
)
- options = []
- if use_tvm_ffi:
- options.append("--enable-tvm-ffi")
-
compiled_fn = compile_bmm(
gemm,
ab_dtype,
@@ -1920,7 +1920,6 @@ def run(
c_major,
max_active_clusters,
current_stream,
- options=",".join(options),
)
# Run and verify BMM with torch
@@ -2043,12 +2042,6 @@ def prepare_parser():
default=False,
help="Use circular buffer tensor sets to ensure L2 cold cache",
)
- parser.add_argument(
- "--use_tvm_ffi",
- action="store_true",
- default=False,
- help="Enable TVM FFI for the kernel, defaults to False using CuTe DSL's native runtime",
- )
return parser
@@ -2090,7 +2083,6 @@ if __name__ == "__main__":
args.iterations,
args.skip_ref_check,
args.use_cold_l2,
- args.use_tvm_ffi,
args.benchmark,
)
print("PASS")
diff --git a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst
index 42d0fdea..2b8a8bc5 100644
--- a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst
+++ b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst
@@ -4,7 +4,8 @@
Compile with TVM FFI
====================
-Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in the `official documentation `_.
+Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in
+the `official documentation `_.
To install TVM FFI, you can run the following command:
@@ -14,7 +15,9 @@ To install TVM FFI, you can run the following command:
# optional package for improved torch tensor calling performance
pip install torch-c-dlpack-ext
-In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster JIT function invocation and provides better interoperability with machine learning frameworks (e.g., directly take ``torch.Tensor`` as arguments).
+In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster
+JIT function invocation and provides better interoperability with machine learning frameworks
+(e.g., directly take ``torch.Tensor`` as arguments).
Enable Apache TVM FFI in |DSL|
@@ -40,7 +43,8 @@ There are two ways to enable TVM FFI in |DSL|:
Note that the object returned by ``cute.compile`` is a Python function specific to TVM FFI.
-2. Alternatively, you can enable TVM FFI globally by setting the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. Please note that this setting will apply to all JIT compilations within the environment.
+2. Alternatively, you can enable TVM FFI globally by setting the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``.
+Please note that this setting will apply to all JIT compilations within the environment.
Minimizing Host Overhead
@@ -129,7 +133,8 @@ stride via the ``stride`` argument in the ``make_fake_tensor`` API.
``cute.Tensor`` adapter for TVM FFI
-----------------------------------
-To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the ``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example:
+To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the
+``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example:
.. code-block:: python
@@ -288,6 +293,33 @@ composed of the types that are supported by TVM FFI. The example below shows how
example_add_one_with_tuple()
+Limitations
+-----------
+
+The Fake Tensor flow is ONLY compatible with TVM FFI because TVM FFI support more flexible constraints on Tensor arguments.
+For instance, fake tensor can specify per-mode static shape or constraints on shape and strides which is not supported by
+existing `from_dlpack` flow. It's expected that JIT function compiled with fake tensor may have different ABI with
+tensor converted by `from_dlpack`.
+
+.. code-block:: python
+
+ import cutlass.cute as cute
+ import torch
+
+ n = cute.sym_int()
+ # Dynamic Shape
+ fake_a = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
+
+ # Compile without tvm-ffi
+ compiled_fn = cute.compile(foo, fake_a)
+
+ # Wrong, in compatible ABI
+ compiled_fn(from_dlpack(a))
+
+
+In order to avoid such issue, it's recommended fake tensor is only used with TVM FFI.
+
+
Supported types
---------------