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

[Ammar-Alnagar]

1. What was the Issue? (#3527)

In CUTLASS v4.6, the routing condition in python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py was updated to check:

if (
    kwargs_wrapper_spec.kwonly_names
    or kwargs_wrapper_spec.arg_defaults
    or kwargs_wrapper_spec.arg_names
    or map_dataclass_to_tuple
):
    return TVMFFIJitCompiledFunctionWithKwargs(...)

Because kwargs_wrapper_spec.arg_names includes positional parameter names for every argument, any @cute.jit function taking 1 or more arguments evaluated to True and was routed to TVMFFIJitCompiledFunctionWithKwargs.

TVMFFIJitCompiledFunctionWithKwargs wraps invocation in kwargs_wrapper.make_kwargs_wrapper. When registered via TVM-FFI and called from C++, execution left C++ to enter Python (GIL + argument wrapper parsing), causing CPU overhead regression compared to direct C++ TVM-FFI function execution.

---

2. What We Did

Positional-only parameters (def func(a, b, c, /)) cannot be called with keyword arguments in Python grammar. Therefore, functions with strictly positional-only arguments do not require keyword-argument wrapping.

1. python/CuTeDSL/cutlass/base_dsl/jit_executor.py:

   • Added has_pos_or_kw to KwargsWrapperSpec.
   • Updated get_kwargs_wrapper_spec to track whether any Parameter.POSITIONAL_OR_KEYWORD parameter exists.

2. python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py:

   • Changed kwargs_wrapper_spec.arg_names check in _make_compiled_func to kwargs_wrapper_spec.has_pos_or_kw.
   • Functions defined as def func(a, b, c, /) without defaults or dataclasses now return TVMFFIJitCompiledFunction (a direct subclass of tvm_ffi.Function).

3. Unit Tests & Documentation:

   • Added unit test suite test_tvm_ffi_kwargs_wrapper_spec.py.
   • Updated documentation in tvm_ffi_compilation.rst:64.

---

3. Comparison & Verification of Fix

Function Signature	has_pos_or_kw	arg_defaults	kwonly_names	Class Returned	CPU Overhead Behavior
def func()	False	()	[]	TVMFFIJitCompiledFunction	0 Python Overhead (Direct C++ tvm_ffi.Function)
def func(a, b, c, /) (Fix)	False	()	[]	TVMFFIJitCompiledFunction	0 Python Overhead (Direct C++ tvm_ffi.Function, bypasses GIL/wrapper)
def func(a, b, c)	True	()	[]	TVMFFIJitCompiledFunctionWithKwargs	Handles Python kwargs calls (func(a=1, b=2))
def func(a, b=1, /)	False	(1,)	[]	TVMFFIJitCompiledFunctionWithKwargs	Handles positional default value binding
def func(a, /, *, k=1)	False	()	['k']	TVMFFIJitCompiledFunctionWithKwargs	Handles keyword-only argument binding

Verification Summary

• Evaluated signature parsing with inspect.signature across positional-only, positional-or-keyword, and mixed argument patterns.
• Confirmed def func(a, b, c, /) correctly evaluates check_routing() to False, bypassing TVMFFIJitCompiledFunctionWithKwargs and directly instantiating TVMFFIJitCompiledFunction.

[Ammar-Alnagar]

Hey @kainzhong / @tqchen,

Following up on the earlier discussion ([https://github.com//issues/3527]):

Hi @kainzhong, thanks for reaching out! This is a known issue for us. Your understanding is correct - the condition you mention is generally true, except the trivial functions without any args. But per our benchmark, the overhead it introduces is < 1us, while the overall launch overhead is typically <10 us. We think it is fine to keep this overhead for compatibility of diversed inputs. And we agree to have more documents here. Any PR is appreciated.
cc: @tqchen

This PR addresses the routing issue so that pure positional-only functions (def func(a, b, c, /)) now correctly bypass TVMFFIJitCompiledFunctionWithKwargs and use the direct TVMFFIJitCompiledFunction path (zero Python/GIL overhead).

Would appreciate a review when you have a chance — thanks!

[kainzhong]

Hi @cyx-6 can you take a look at this PR? My opinion is

• Should we make this behavior even more explicit? Now adding / can change the CPU overhead feels a bit black magic and I wonder if it would be better to add a param to cute.compile like optimize_for_pos_only
• I don't think test/python/CuTeDSL/test_tvm_ffi_cpp_registry.py is needed
• test_kwargs_wrapper_spec_positional_only, test_kwargs_wrapper_spec_positional_or_keyword and test_kwargs_wrapper_spec_mixed_posonly_and_pos_or_kw seem unnecessary as well

But I think that should be up to cutlass people to decide. What do you think?

[cyx-6]

@kainzhong Thanks for contribution! This change looks good to me. And yes, I agree that some semantic changes should be discussed and decided by cutlass people.

[Ammar-Alnagar]

Hi @cyx-6 can you take a look at this PR? My opinion is

• Should we make this behavior even more explicit? Now adding / can change the CPU overhead feels a bit black magic and I wonder if it would be better to add a param to cute.compile like optimize_for_pos_only
• I don't think test/python/CuTeDSL/test_tvm_ffi_cpp_registry.py is needed
• test_kwargs_wrapper_spec_positional_only, test_kwargs_wrapper_spec_positional_or_keyword and test_kwargs_wrapper_spec_mixed_posonly_and_pos_or_kw seem unnecessary as well

But I think that should be up to cutlass people to decide. What do you think?

I can see where your point in the black box change and an input from one of the cutlass people would be much appreciated but i think the test cases don't hurt and are more of an assurance , but thats just my opinion and might not be the best flow for the project .
Curious on what you think? @cyx-6 @kainzhong

[kainzhong]

Hi @Ammar-Alnagar unfortunately I'm not a cutlass developer. I think you probably should request their review instead since they are the ones who can approve it.

[Ammar-Alnagar]

Hi @Ammar-Alnagar unfortunately I'm not a cutlass developer. I think you probably should request their review instead since they are the ones who can approve it.

no worries , thanks alot for the help and comments !!!

[Ammar-Alnagar]

@Junkai-Wu @brandon-yujie-sun @fengxie , can i get a review on this whenever you get a chance ? , thanks!!!

[brandon-yujie-sun]

Thanks for contributing the change!

[brandon-yujie-sun]

LGTM, thanks

[brandon-yujie-sun]

@Junkai-Wu for further processing