Skip to content

fix(CuTeDSL): preserve wrapped signatures for type checkers - #3528

Open
iwzbi wants to merge 1 commit into
NVIDIA:mainfrom
iwzbi:fix/preserve-dsl-user-op-signatures
Open

fix(CuTeDSL): preserve wrapped signatures for type checkers#3528
iwzbi wants to merge 1 commit into
NVIDIA:mainfrom
iwzbi:fix/preserve-dsl-user-op-signatures

Conversation

@iwzbi

@iwzbi iwzbi commented Aug 25, 2026

Copy link
Copy Markdown

Summary

@dsl_user_op decorates every user-facing CuTe DSL wrapper (local_tile, make_tiled_mma, copy, gemm, etc.), but its annotation was Callable[..., Any] -> Callable[..., Any], which erases the wrapped callable's signature for static analyzers.

Pyright/basedpyright (and language servers built on them) therefore see only (*args, **kwargs). In practice this means:

  • Editors show no parameter names in signature help: cute.local_tile( yields no tiler/coord hints;
  • Bogus keyword calls compile silently (e.g. cute.local_tile(mA, xxxxx=...) passes type checking);
  • Downstream projects that write their own @dsl_user_op wrappers (vLLM, SGLang, FlashInfer) are affected the same way.

Change

Make the decorator signature-preserving with a TypeVar bound to Callable, the pattern recommended by pyright's documentation:

_FuncT = TypeVar("_FuncT", bound=Callable[..., Any])

def dsl_user_op(opFunc: _FuncT) -> _FuncT:
    ...
    return wrapper  # type: ignore[return-value]

Runtime behavior is unchanged (functools.wraps was already applied to the wrapper). # type: ignore[return-value] documents that the inner wrapper intentionally has a generic signature but must satisfy _FuncT.

Verification

  • python -m py_compile op.py — syntax OK.
  • pyright, after patch: cute.local_tile(mA, tiler=(128, 64), coord=(0, 0)) type-checks; cute.local_tile(mA, xxxxx=(128, 64)) is rejected with No parameter named "xxxxx", and Arguments missing for parameters "tiler", "coord" is reported for incomplete calls.

Fixes editor/signature-help blindness for the entire public CuTe DSL surface.

dsl_user_op is applied to every user-facing CuTe DSL wrapper (local_tile,
make_tiled_mma, gemm, etc.) but its annotation erased the wrapped callable
to Callable[..., Any]. Static analyzers (pyright/basedpyright/mypy) and
language servers therefore see only (*args, **kwargs): editors show no
parameter names in signature help and IDEs report bogus unknown-keyword
calls as valid.

Use a TypeVar bound to Callable so the decorator returns the exact
signature of the wrapped function, following the recommended decorator
pattern in the pyright documentation:

    _FuncT = TypeVar('_FuncT', bound=Callable[..., Any])
    def decorator(func: _FuncT) -> _FuncT: ...

Verified with pyright: cute.local_tile(mA, tiler=..., coord=...) now
resolves its real named parameters, while cute.local_tile(mA, xxxxx=...)
is correctly rejected as 'No parameter named "xxxxx"'.

Runtime behavior is unchanged (functools.wraps was already applied).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant