From 6a1e3f668995c9f3dce937a68aaf245c546bc03a Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 17 Apr 2026 09:17:39 +0000 Subject: [PATCH 1/3] [None][perf] reduce @torch.library.custom_op host overhead Switch trtllm::nvfp4_gemm and trtllm::tunable_fp4_quantize from the @torch.library.custom_op decorator to the low-level torch.library.Library.define + impl API. The high-level decorator carries a hidden ~12us per-call dispatcher tax (visible on Python-heavy, host-bound iterations); the low-level API avoids it while preserving torch.compile support via register_fake. LTX2 dense transformer issues ~1260 tunable_fp4_quantize and ~840 nvfp4_gemm calls per step. At ~12us/call that is ~15ms/step and ~10ms/step of pure CPU dispatcher cost respectively, which amplifies on multi-GPU due to NCCL synchronization. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/custom_ops/torch_custom_ops.py | 42 ++++++++++++++++--- 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 61b466a76fc1..75deab8d47d4 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -904,7 +904,16 @@ def forward( raise ValueError(f"Invalid tactic: {tactic}") -@torch.library.custom_op("trtllm::nvfp4_gemm", mutates_args=()) +# Use low-level torch.library.Library API instead of @torch.library.custom_op to +# avoid the ~12us/call dispatcher tax of the high-level decorator. See PR notes +# for benchmarks. +_trtllm_nvfp4_gemm_lib = torch.library.Library("trtllm", "FRAGMENT") +_trtllm_nvfp4_gemm_lib.define( + "nvfp4_gemm(Tensor act_fp4, Tensor weight, Tensor act_sf, Tensor weight_scale, " + "Tensor alpha, ScalarType output_dtype, bool to_userbuffers=False, " + "str allowed_backends=\"cutlass,cublaslt,cuda_core\") -> Tensor") + + def nvfp4_gemm( act_fp4: torch.Tensor, weight: torch.Tensor, @@ -997,8 +1006,10 @@ def nvfp4_gemm( ) -@nvfp4_gemm.register_fake -def _( +_trtllm_nvfp4_gemm_lib.impl("nvfp4_gemm", nvfp4_gemm, "CUDA") + + +def _nvfp4_gemm_fake( act_fp4: torch.Tensor, weight: torch.Tensor, act_sf: torch.Tensor, @@ -1013,6 +1024,9 @@ def _( dtype=output_dtype) +torch.library.register_fake("trtllm::nvfp4_gemm", _nvfp4_gemm_fake) + + class FP8BatchedGemmRunner(TunableRunner): runner_dict = dict() tuning_config = None @@ -2350,7 +2364,16 @@ def forward( return act_fp4 -@torch.library.custom_op("trtllm::tunable_fp4_quantize", mutates_args=()) +# Use low-level torch.library.Library API instead of @torch.library.custom_op to +# avoid the ~12us/call dispatcher tax of the high-level decorator. The trtllm +# namespace is already defined by C++ via TORCH_LIBRARY_FRAGMENT, so we add to +# it here using FRAGMENT mode. +_trtllm_tunable_fp4_quantize_lib = torch.library.Library("trtllm", "FRAGMENT") +_trtllm_tunable_fp4_quantize_lib.define( + "tunable_fp4_quantize(Tensor input, Tensor input_scale, " + "int scaling_vector_size, bool is_sf_swizzled_layout) -> Tensor[]") + + def tunable_fp4_quantize( input: torch.Tensor, input_scale: torch.Tensor, @@ -2402,8 +2425,11 @@ def tunable_fp4_quantize( return [act_fp4, act_sf] -@tunable_fp4_quantize.register_fake -def _( +_trtllm_tunable_fp4_quantize_lib.impl("tunable_fp4_quantize", + tunable_fp4_quantize, "CUDA") + + +def _tunable_fp4_quantize_fake( input: torch.Tensor, input_scale: torch.Tensor, scaling_vector_size: int = 16, @@ -2424,3 +2450,7 @@ def _( input.new_empty(output_shape, dtype=torch.uint8), input_scale.new_empty(scale_shape, dtype=torch.uint8), ] + + +torch.library.register_fake("trtllm::tunable_fp4_quantize", + _tunable_fp4_quantize_fake) From 012a410261c3bfd25987a1f587da19e01a063982 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 22 Apr 2026 03:30:26 +0000 Subject: [PATCH 2/3] [None][perf] add @fast_custom_op helper preserving @custom_op ergonomics Extract the low-level torch.library.Library.define+impl pattern into a @fast_custom_op decorator that keeps the @torch.library.custom_op developer experience (schema inferred from Python type hints via infer_schema, same .register_fake method on the returned op) while bypassing the ~7us/call Python dispatcher tax. Switch trtllm::nvfp4_gemm and trtllm::tunable_fp4_quantize from the manual Library.define+impl form to @fast_custom_op. This: - Keeps type-hint-driven schema inference (no hand-written schema strings) - Restores @op.register_fake ergonomics - Serves as a template for migrating more @custom_op sites mechanically Microbenchmark (B200, PyTorch 2.10, 20k-iter tight loop, x.clone() kernel): @torch.library.custom_op 11.67us/call (7.02us dispatcher tax) Manual Library.define+impl 6.28us/call (1.63us tax, -5.39us) @fast_custom_op via torch.ops 6.28us/call (1.63us tax, -5.38us) @fast_custom_op via proxy call 6.09us/call (1.44us tax, -5.57us) The helper is zero-cost on the hot path (torch.ops.trtllm.(...) goes through the C++ dispatcher directly, same as the manual form). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/custom_ops/fast_custom_op.py | 118 ++++++++++++++++++ .../_torch/custom_ops/torch_custom_ops.py | 39 +----- 2 files changed, 123 insertions(+), 34 deletions(-) create mode 100644 tensorrt_llm/_torch/custom_ops/fast_custom_op.py diff --git a/tensorrt_llm/_torch/custom_ops/fast_custom_op.py b/tensorrt_llm/_torch/custom_ops/fast_custom_op.py new file mode 100644 index 000000000000..8476347618b3 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/fast_custom_op.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Low-overhead replacement for ``@torch.library.custom_op``. + +``@torch.library.custom_op`` is ergonomic (auto schema inference from Python +type hints, a ``register_fake`` method on the returned op) but its wrapper +imposes a ~6-7us per-call dispatcher tax (schema re-validation, Python-level +``DispatchKeySet`` traversal, auto-functionalization bookkeeping, etc.). + +``fast_custom_op`` preserves the ergonomics while bypassing that tax by +registering the op directly through the low-level +``torch.library.Library.define + impl`` API — same path as built-in ATen ops. + +Usage is almost identical to ``@torch.library.custom_op``:: + + from tensorrt_llm._torch.custom_ops.fast_custom_op import fast_custom_op + + @fast_custom_op("trtllm::nvfp4_gemm", mutates_args=()) + def nvfp4_gemm(x: torch.Tensor, ...) -> torch.Tensor: + ... + + @nvfp4_gemm.register_fake + def _(x: torch.Tensor, ...) -> torch.Tensor: + ... + +Caveats (inherited from using the low-level API): + +* No autograd support (register a separate autograd kernel if the op is + differentiable — but in practice this is for stateless inference kernels). +* ``mutates_args`` must be a concrete tuple; ``"unknown"`` auto-functionalization + is not supported. +* ``device_types`` defaults to ``"CUDA"``. Pass a different string or a tuple + of strings to register on other backends. +""" + +from __future__ import annotations + +from typing import Callable, Iterable, Tuple, Union + +import torch +from torch.library import Library, infer_schema, register_fake + +_LIBS: dict[tuple[str, str], Library] = {} + + +def _get_library(namespace: str, kind: str = "FRAGMENT") -> Library: + key = (namespace, kind) + lib = _LIBS.get(key) + if lib is None: + lib = Library(namespace, kind) + _LIBS[key] = lib + return lib + + +def fast_custom_op( + qualname: str, + *, + mutates_args: Union[Iterable[str], str] = (), + device_types: Union[str, Tuple[str, ...]] = "CUDA", +) -> Callable[[Callable], "FastCustomOp"]: + """Register a Python function as a fast custom torch op. + + Parameters mirror ``torch.library.custom_op``: + qualname: ``"::"`` identifier. + mutates_args: names of arguments that are mutated in-place; empty tuple + means the op is pure. + device_types: backend(s) to register the impl on (default ``"CUDA"``). + """ + if "::" not in qualname: + raise ValueError(f"qualname must be '::', got {qualname!r}") + namespace, op_name = qualname.split("::", 1) + + if isinstance(mutates_args, str) and mutates_args != "unknown": + raise TypeError("mutates_args must be an iterable of names or 'unknown'") + mutates_args_tuple = mutates_args if isinstance(mutates_args, str) else tuple(mutates_args) + + dev_types = (device_types,) if isinstance(device_types, str) else tuple(device_types) + + def decorator(fn: Callable) -> "FastCustomOp": + schema = infer_schema(fn, op_name=op_name, mutates_args=mutates_args_tuple) + lib = _get_library(namespace) + lib.define(schema) + for dt in dev_types: + lib.impl(op_name, fn, dt) + return FastCustomOp(qualname=qualname, namespace=namespace, op_name=op_name, python_fn=fn) + + return decorator + + +class FastCustomOp: + """Handle returned by :func:`fast_custom_op`. + + Behaves like ``@torch.library.custom_op``'s return value: callable and + exposes ``register_fake``. The call path goes through the C++ dispatcher + (``torch.ops..``), bypassing the Python wrapper layer of + ``@custom_op``. + """ + + __slots__ = ("qualname", "namespace", "op_name", "_python_fn", "_op") + + def __init__(self, qualname: str, namespace: str, op_name: str, python_fn: Callable): + self.qualname = qualname + self.namespace = namespace + self.op_name = op_name + self._python_fn = python_fn + self._op = getattr(getattr(torch.ops, namespace), op_name) + + def __call__(self, *args, **kwargs): + return self._op(*args, **kwargs) + + def register_fake(self, fake_fn: Callable) -> Callable: + register_fake(self.qualname, fake_fn) + return fake_fn + + @property + def python_impl(self) -> Callable: + """The original un-wrapped Python function (for tests/introspection).""" + return self._python_fn diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 75deab8d47d4..0c1cd250491d 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -20,6 +20,7 @@ from ..cublaslt_utils import IS_CUBLASLT_AVAILABLE from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from .fast_custom_op import fast_custom_op if IS_FLASHINFER_AVAILABLE: from flashinfer.fp4_quantization import nvfp4_quantize as _flashinfer_nvfp4_quantize @@ -904,16 +905,7 @@ def forward( raise ValueError(f"Invalid tactic: {tactic}") -# Use low-level torch.library.Library API instead of @torch.library.custom_op to -# avoid the ~12us/call dispatcher tax of the high-level decorator. See PR notes -# for benchmarks. -_trtllm_nvfp4_gemm_lib = torch.library.Library("trtllm", "FRAGMENT") -_trtllm_nvfp4_gemm_lib.define( - "nvfp4_gemm(Tensor act_fp4, Tensor weight, Tensor act_sf, Tensor weight_scale, " - "Tensor alpha, ScalarType output_dtype, bool to_userbuffers=False, " - "str allowed_backends=\"cutlass,cublaslt,cuda_core\") -> Tensor") - - +@fast_custom_op("trtllm::nvfp4_gemm", mutates_args=()) def nvfp4_gemm( act_fp4: torch.Tensor, weight: torch.Tensor, @@ -1006,9 +998,7 @@ def nvfp4_gemm( ) -_trtllm_nvfp4_gemm_lib.impl("nvfp4_gemm", nvfp4_gemm, "CUDA") - - +@nvfp4_gemm.register_fake def _nvfp4_gemm_fake( act_fp4: torch.Tensor, weight: torch.Tensor, @@ -1024,9 +1014,6 @@ def _nvfp4_gemm_fake( dtype=output_dtype) -torch.library.register_fake("trtllm::nvfp4_gemm", _nvfp4_gemm_fake) - - class FP8BatchedGemmRunner(TunableRunner): runner_dict = dict() tuning_config = None @@ -2364,16 +2351,7 @@ def forward( return act_fp4 -# Use low-level torch.library.Library API instead of @torch.library.custom_op to -# avoid the ~12us/call dispatcher tax of the high-level decorator. The trtllm -# namespace is already defined by C++ via TORCH_LIBRARY_FRAGMENT, so we add to -# it here using FRAGMENT mode. -_trtllm_tunable_fp4_quantize_lib = torch.library.Library("trtllm", "FRAGMENT") -_trtllm_tunable_fp4_quantize_lib.define( - "tunable_fp4_quantize(Tensor input, Tensor input_scale, " - "int scaling_vector_size, bool is_sf_swizzled_layout) -> Tensor[]") - - +@fast_custom_op("trtllm::tunable_fp4_quantize", mutates_args=()) def tunable_fp4_quantize( input: torch.Tensor, input_scale: torch.Tensor, @@ -2425,10 +2403,7 @@ def tunable_fp4_quantize( return [act_fp4, act_sf] -_trtllm_tunable_fp4_quantize_lib.impl("tunable_fp4_quantize", - tunable_fp4_quantize, "CUDA") - - +@tunable_fp4_quantize.register_fake def _tunable_fp4_quantize_fake( input: torch.Tensor, input_scale: torch.Tensor, @@ -2450,7 +2425,3 @@ def _tunable_fp4_quantize_fake( input.new_empty(output_shape, dtype=torch.uint8), input_scale.new_empty(scale_shape, dtype=torch.uint8), ] - - -torch.library.register_fake("trtllm::tunable_fp4_quantize", - _tunable_fp4_quantize_fake) From bc098214a4bf8bb2e2170e72b2faa232db54c6d9 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 22 Apr 2026 06:22:34 +0000 Subject: [PATCH 3/3] [None][perf] minimize diff: revert fake function names to _ Revert fake function names back to _ to match the rest of the file (e.g. nvfp4_gemm_cublaslt, fp8_rowwise_gemm) and the original @custom_op idiom. The named identifiers were only needed during the intermediate manual Library.define+impl step; with @fast_custom_op the decorator form no longer needs named identifiers. No behavior change. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 0c1cd250491d..90cc4ecdc582 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -999,7 +999,7 @@ def nvfp4_gemm( @nvfp4_gemm.register_fake -def _nvfp4_gemm_fake( +def _( act_fp4: torch.Tensor, weight: torch.Tensor, act_sf: torch.Tensor, @@ -2404,7 +2404,7 @@ def tunable_fp4_quantize( @tunable_fp4_quantize.register_fake -def _tunable_fp4_quantize_fake( +def _( input: torch.Tensor, input_scale: torch.Tensor, scaling_vector_size: int = 16,