fix(CuTeDSL): optimize positional-only functions with TVMFFIJitCompiledFunction to eliminate CPU overhead (#3527)

media/docs/pythonDSL/guides/tvm_ffi_compilation.rst

@@ -61,12 +61,39 @@ To maximize performance benefits, we recommend setting up your workflow as follo
 - **Use the environment stream flag** to implicitly pass the current PyTorch stream.
 - **Rely on compiled argument validation** instead of Python-side attribute validation,
   as TVM FFI functions perform fast compiled checks.
+- **Use the positional-only TVM FFI fast path** when its signature restrictions fit your API.
 
 Following these steps can significantly reduce the host-side overhead of eager kernel execution.
 The sections below provide detailed examples and explanations for each step.
 You may find it helpful to refer back to this summary after you review the implementation details.
 
 
+Use the positional-only TVM FFI fast path
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+TVM FFI's native ABI accepts positional arguments only. CuTe DSL therefore
+returns a direct ``tvm_ffi.Function`` for a ``@cute.jit`` function whose
+runtime arguments are all positional-only and have no defaults. This avoids
+the Python keyword/defaults wrapper and its GIL overhead for both Python and
+C++ callers.
+
+For example, this signature uses the direct path:
+
+.. code-block:: python
+
+   @cute.jit
+   def add_one(a: cute.Tensor, b: cute.Tensor, /):
+      ...
+
+The direct path cannot accept keyword arguments. CuTe DSL uses
+``TVMFFIJitCompiledFunctionWithKwargs`` instead when a signature has a
+positional-or-keyword parameter, a default value, a keyword-only parameter,
+or a top-level dataclass argument. The wrapper preserves Python calls such as
+``compiled_add_one(a=a_torch, b=b_torch)`` and default-value handling, but has
+Python wrapper overhead. In particular, a positional-only signature with a
+default (``def func(a, b=1, /)``) still uses the wrapper.
+
+
 Fake tensor for compilation
 ---------------------------
 
@@ -90,7 +117,7 @@ To clearly distinguish between the compilation phase and runtime,
          b[tid] = a[tid] + 1.0
 
    @cute.jit
-   def add_one(a: cute.Tensor, b: cute.Tensor):
+   def add_one(a: cute.Tensor, b: cute.Tensor, /):
       n = a.shape[0]
       threads_per_block = 128
       blocks = (n + threads_per_block - 1) // threads_per_block

python/CuTeDSL/cutlass/base_dsl/jit_executor.py

@@ -215,6 +215,7 @@ class KwargsWrapperSpec(NamedTuple):
     arg_defaults: tuple[Any, ...]
     kwonly_names: list[str]
     kwonly_defaults: dict[str, Any]
+    has_pos_or_kw: bool = True
 
 
 def _validate_pointer_address(address: int) -> None:
@@ -821,6 +822,7 @@ def get_kwargs_wrapper_spec(
         arg_defaults = []
         kwonly_names = []
         kwonly_defaults = {}
+        has_pos_or_kw = False
 
         for i, (name, param) in enumerate(sig.parameters.items()):
             # We don't want *args or **kwargs to be included in the
@@ -831,13 +833,16 @@ def get_kwargs_wrapper_spec(
                 inspect.Parameter.VAR_KEYWORD,
             ):
                 continue
-            is_kwonly = param.kind == inspect.Parameter.KEYWORD_ONLY
             annotation = param.annotation
             if (
                 is_arg_annotation_constexpr(annotation, name, i, None)
                 or name in excluded_arg_names
             ):
                 continue
+            is_kwonly = param.kind == inspect.Parameter.KEYWORD_ONLY
+            is_pos_or_kw = param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
+            if is_pos_or_kw:
+                has_pos_or_kw = True
             arg_names.append(name) if not is_kwonly else kwonly_names.append(name)
             if param.default is not inspect.Parameter.empty:
                 if is_kwonly:
@@ -850,6 +855,7 @@ def get_kwargs_wrapper_spec(
             arg_defaults=tuple(arg_defaults),
             kwonly_names=kwonly_names,
             kwonly_defaults=kwonly_defaults,
+            has_pos_or_kw=has_pos_or_kw,
         )
 
     def get_rectified_args_from_original_args(

python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py

@@ -1016,15 +1016,15 @@ def post_compile_hook(module: ir.Module) -> None:
 
             def _make_compiled_func(*args: Any, **kwargs: Any) -> Any:
                 # Route through the kwargs-capable compiled class whenever
-                # the signature has kwonly/defaults OR any positional arg
-                # carries a dataclass instance (detected by the spec
-                # converter, including Union[...] over dataclasses). The
-                # kwargs wrapper is the only place tvm-ffi's
+                # the signature permits keyword arguments, has defaults, or
+                # any positional arg carries a dataclass instance (detected
+                # by the spec converter, including Union[...] over
+                # dataclasses). The kwargs wrapper is the only place tvm-ffi's
                 # ``map_dataclass_to_tuple`` unpack hook fires.
                 if (
                     kwargs_wrapper_spec.kwonly_names
                     or kwargs_wrapper_spec.arg_defaults
-                    or kwargs_wrapper_spec.arg_names
+                    or kwargs_wrapper_spec.has_pos_or_kw
                     or map_dataclass_to_tuple
                 ):
                     return TVMFFIJitCompiledFunctionWithKwargs(

test/python/CuTeDSL/test_tvm_ffi_kwargs_wrapper_spec.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import inspect
+import pytest
+
+pytest.importorskip("tvm_ffi")
+
+import cutlass.cute as cute
+from cutlass import Int32
+from cutlass.base_dsl.jit_executor import ExecutionArgs
+from cutlass.cutlass_dsl.tvm_ffi_provider import (
+    TVMFFIJitCompiledFunction,
+    TVMFFIJitCompiledFunctionWithKwargs,
+)
+
+
+@cute.jit
+def _positional_only(a: Int32, b: Int32, /):
+    pass
+
+
+@cute.jit
+def _positional_or_keyword(a: Int32, b: Int32):
+    pass
+
+
+@cute.jit
+def _positional_default(a: Int32, b: Int32 = Int32(2), /):
+    pass
+
+
+@cute.jit
+def _keyword_only(a: Int32, /, *, k: Int32 = Int32(2)):
+    pass
+
+
+def _compile(func, *args):
+    return cute.compile(func, *args, options="--enable-tvm-ffi")
+
+
+def test_tvm_ffi_routes_positional_only_signature_to_native_function():
+    compiled = _compile(_positional_only, Int32(1), Int32(2))
+
+    assert isinstance(compiled, TVMFFIJitCompiledFunction)
+    compiled(1, 2)
+
+
+def test_tvm_ffi_routes_positional_or_keyword_signature_to_kwargs_wrapper():
+    compiled = _compile(_positional_or_keyword, Int32(1), Int32(2))
+
+    assert isinstance(compiled, TVMFFIJitCompiledFunctionWithKwargs)
+    compiled(a=1, b=2)
+
+
+def test_tvm_ffi_routes_positional_default_signature_to_kwargs_wrapper():
+    compiled = _compile(_positional_default, Int32(1))
+
+    assert isinstance(compiled, TVMFFIJitCompiledFunctionWithKwargs)
+    compiled(1)
+
+
+def test_tvm_ffi_routes_keyword_only_signature_to_kwargs_wrapper():
+    compiled = _compile(_keyword_only, Int32(1))
+
+    assert isinstance(compiled, TVMFFIJitCompiledFunctionWithKwargs)
+    compiled(1, k=2)
+
+
+def test_kwargs_wrapper_spec_positional_only():
+    def func_posonly(a, b, c, /):
+        pass
+
+    sig = inspect.signature(func_posonly)
+    exec_args = ExecutionArgs(sig, "func_posonly")
+    spec = exec_args.get_kwargs_wrapper_spec()
+
+    assert spec.arg_names == ["a", "b", "c"]
+    assert spec.arg_defaults == ()
+    assert spec.kwonly_names == []
+    assert spec.kwonly_defaults == {}
+    assert spec.has_pos_or_kw is False
+
+
+def test_kwargs_wrapper_spec_positional_or_keyword():
+    def func_pos_or_kw(a, b, c):
+        pass
+
+    sig = inspect.signature(func_pos_or_kw)
+    exec_args = ExecutionArgs(sig, "func_pos_or_kw")
+    spec = exec_args.get_kwargs_wrapper_spec()
+
+    assert spec.arg_names == ["a", "b", "c"]
+    assert spec.has_pos_or_kw is True
+
+
+def test_kwargs_wrapper_spec_mixed_posonly_and_pos_or_kw():
+    def func_mixed(a, /, b):
+        pass
+
+    sig = inspect.signature(func_mixed)
+    exec_args = ExecutionArgs(sig, "func_mixed")
+    spec = exec_args.get_kwargs_wrapper_spec()
+
+    assert spec.arg_names == ["a", "b"]
+    assert spec.has_pos_or_kw is True