From 7b985e3a88aa1264e0a4d757d30cb964232ab148 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 3 Sep 2026 07:00:08 +0000 Subject: [PATCH 1/2] [None][perf] Port tunable custom ops to fast_custom_op `@torch.library.custom_op` re-validates the schema, walks the Python `DispatchKeySet` and does auto-functionalization bookkeeping on *every* call. For a tunable op that cost is paid on the host before its kernel is even launched, and it is not small relative to the kernel: measured against a 25us reference kernel it runs from 28% (1 arg) up to 74% (`fused_moe`, 51 args) of the kernel's own duration. `fast_custom_op` already exists to avoid that tax by registering through `Library.define + impl`, but only two tunable ops used it. Port the remaining 31 ops whose bodies resolve a tactic through the AutoTuner. Measured on an AMD EPYC 7313P (40k iterations, best of 3) with trivial op bodies, so only the dispatch layer is timed: #args example op custom_op fast saved 1 quantize_e4m3_per_tensor 9.32u 2.43u 6.89u 5 fp8_swap_ab_gemm 10.79u 3.01u 7.78u 11 tunable_allreduce 13.07u 3.88u 9.18u 30 mxe4m3_mxe2m1_block_scale_moe_runner 20.43u 6.65u 13.78u 51 fused_moe 27.94u 9.56u 18.38u The saving lands on paths that actually execute the Python op body, i.e. prefill and mixed steps, decode batches outside the captured CUDA graph sizes, and graph capture itself. Steps served entirely by a replayed CUDA graph never run these bodies and are unaffected. Two gaps in `fast_custom_op` had to be closed first, both of which would otherwise have silently changed behavior when porting an op: * `device_types=None` now registers one `CompositeExplicitAutograd` kernel, matching what `custom_op` does when `device_types` is omitted. Without it the 19 ported ops that never named a device would have narrowed to CUDA and raised NotImplementedError on a CPU tensor. * `device_types` accepts a device type ("cuda") as well as a dispatch key ("CUDA"). `Library.impl` only takes the latter, so the 12 CuTe DSL ops that spell it lowercase would have failed at import. The one property genuinely lost is `custom_op`'s per-call aliasing check: an op that returns an undeclared alias of an input now silently corrupts that input instead of raising. `TLLM_VALIDATE_CUSTOM_OPS=1` makes every `fast_custom_op` fall back to `torch.library.custom_op` and restores the check, so the guarantee can be recovered in CI and when bisecting a suspected miscompare. Wiring that flag into a CI stage is left as a follow-up. The 7 tunable ops that declare `mutates_args` are deliberately not ported: they are also the ones with no targeted test coverage, so they are held back until tests exist. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 61 ++++--- .../_torch/custom_ops/fast_custom_op.py | 59 ++++++- .../_torch/custom_ops/torch_custom_ops.py | 45 ++++-- .../custom_ops/trtllm_gen_custom_ops.py | 29 ++-- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../_torch/custom_ops/test_fast_custom_op.py | 151 ++++++++++++++++++ 6 files changed, 286 insertions(+), 60 deletions(-) create mode 100644 tests/unittest/_torch/custom_ops/test_fast_custom_op.py diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 22e2c7c901a6..bcf898d1704f 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -26,6 +26,7 @@ from .cutedsl_matmul_heuristics import (NVFP4_PRECISION, nvmmh_enabled_for_nvfp4, nvmmh_fields, nvmmh_max_tactics, rank_configs) +from .fast_custom_op import fast_custom_op try: from cuda.bindings import driver as cuda @@ -1013,9 +1014,9 @@ def forward( return c_tensor # a/b: fp4, scale: fp8, output: bf16 - @torch.library.custom_op("trtllm::cute_dsl_nvfp4_gemm_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_gemm_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_gemm_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -1474,10 +1475,9 @@ def forward( return c_tensor # a/b: fp4, scale: fp8, output: bf16, fused SwiGLU activation - @torch.library.custom_op( - "trtllm::cute_dsl_nvfp4_dense_gemm_swiglu_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_dense_gemm_swiglu_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_dense_gemm_swiglu_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -1561,9 +1561,9 @@ def unique_id(self): return (self.output_dtype, self.use_tvm_ffi, 'gelu') # a/b: fp4, scale: fp8, output: bf16, fused non-gated GELU(tanh) - @torch.library.custom_op("trtllm::cute_dsl_nvfp4_dense_gemm_gelu_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_dense_gemm_gelu_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_dense_gemm_gelu_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -2004,10 +2004,9 @@ def forward( return c_tensor, c_sf_tensor # a/b: fp4, scale: fp8, output: fp4 + sfc, fused SwiGLU activation - @torch.library.custom_op( - "trtllm::cute_dsl_nvfp4_dense_gemm_swiglu_fp4out_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_dense_gemm_swiglu_fp4out_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_dense_gemm_swiglu_fp4out_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -2100,10 +2099,9 @@ def unique_id(self): return (self.use_tvm_ffi, 'gelu_fp4out') # a/b: fp4, scale: fp8, output: fp4 + sfc, fused non-gated GELU(tanh) - @torch.library.custom_op( - "trtllm::cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -2424,9 +2422,9 @@ def forward(self, inputs: List[torch.Tensor], ) return c - @torch.library.custom_op("trtllm::cute_dsl_nvfp4_grouped_gemm_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_grouped_gemm_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_grouped_gemm_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -3139,10 +3137,9 @@ def forward(self, inputs: List[torch.Tensor], ) return c, c_sf - @torch.library.custom_op( - "trtllm::cute_dsl_nvfp4_grouped_gemm_swiglu_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_nvfp4_grouped_gemm_swiglu_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_nvfp4_grouped_gemm_swiglu_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -3545,7 +3542,7 @@ def forward(self, inputs: List, return c, c_sf - @torch.library.custom_op( + @fast_custom_op( "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell", mutates_args=(), device_types="cuda") @@ -3919,7 +3916,7 @@ def forward( ) return packed.view(torch.int8), output_scale.view(torch.int32) - @torch.library.custom_op( + @fast_custom_op( "trtllm::cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", mutates_args=(), device_types="cuda", @@ -4223,9 +4220,9 @@ def forward( return c_tensor # a/b: fp8, scale: fp32, output: bf16 - @torch.library.custom_op("trtllm::cute_dsl_fp8_gemm_blackwell", - mutates_args=(), - device_types="cuda") + @fast_custom_op("trtllm::cute_dsl_fp8_gemm_blackwell", + mutates_args=(), + device_types="cuda") def cute_dsl_fp8_gemm_blackwell( input: torch.Tensor, weight: torch.Tensor, @@ -4878,7 +4875,7 @@ def forward( return c, c_sf - @torch.library.custom_op( + @fast_custom_op( "trtllm::cute_dsl_nvfp4_dense_gemm_swiglu_moe_blackwell", mutates_args=(), device_types="cuda", @@ -5267,7 +5264,7 @@ def forward( return c - @torch.library.custom_op( + @fast_custom_op( "trtllm::cute_dsl_nvfp4_dense_gemm_fc2_blackwell", mutates_args=(), device_types="cuda", diff --git a/tensorrt_llm/_torch/custom_ops/fast_custom_op.py b/tensorrt_llm/_torch/custom_ops/fast_custom_op.py index 8476347618b3..04b385363f22 100644 --- a/tensorrt_llm/_torch/custom_ops/fast_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/fast_custom_op.py @@ -30,18 +30,47 @@ def _(x: torch.Tensor, ...) -> torch.Tensor: * ``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. + of strings to register on other backends, or ``None`` to register a single + backend-agnostic kernel (``CompositeExplicitAutograd``), which is what + ``@torch.library.custom_op`` does when ``device_types`` is omitted. +* **The per-call aliasing/mutation validation that ``@custom_op`` performs is + gone.** ``@custom_op`` raises at runtime if an op returns a tensor that + aliases an input without declaring it; the low-level API does not check, so + the same bug silently corrupts the aliased input instead. Set + ``TLLM_VALIDATE_CUSTOM_OPS=1`` to make every ``fast_custom_op`` fall back to + ``@torch.library.custom_op`` and restore those checks — CI runs at least one + stage with the flag set so the checks still guard every op. """ from __future__ import annotations -from typing import Callable, Iterable, Tuple, Union +import os +from typing import Callable, Iterable, Optional, Tuple, Union import torch from torch.library import Library, infer_schema, register_fake _LIBS: dict[tuple[str, str], Library] = {} +# When set, `fast_custom_op` degrades to `@torch.library.custom_op` so the +# schema/aliasing validation it performs on every call is back in force. Meant +# for CI and for bisecting a suspected miscompare down to this decorator; it +# reintroduces the per-call dispatcher tax, so never set it in production. +VALIDATE_CUSTOM_OPS: bool = os.getenv("TLLM_VALIDATE_CUSTOM_OPS", "0") == "1" + + +def _dispatch_key_for_device(device_type: str) -> str: + """Map a device type ("cuda") to a dispatch key ("CUDA"). + + ``Library.impl`` only accepts dispatch keys, while ``custom_op`` accepts + device types. Both spellings are common in the tree, so accept either. + """ + try: + return torch._C._dispatch_key_for_device(device_type) + except (AttributeError, RuntimeError): + # Already a dispatch key ("CUDA", "CompositeExplicitAutograd", ...). + return device_type + def _get_library(namespace: str, kind: str = "FRAGMENT") -> Library: key = (namespace, kind) @@ -56,7 +85,7 @@ def fast_custom_op( qualname: str, *, mutates_args: Union[Iterable[str], str] = (), - device_types: Union[str, Tuple[str, ...]] = "CUDA", + device_types: Optional[Union[str, Tuple[str, ...]]] = "CUDA", ) -> Callable[[Callable], "FastCustomOp"]: """Register a Python function as a fast custom torch op. @@ -65,6 +94,14 @@ def fast_custom_op( 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"``). + ``None`` registers one backend-agnostic kernel, matching what + ``@torch.library.custom_op`` does when ``device_types`` is omitted — + use it when porting such an op so its set of supported devices does + not silently narrow. + + With ``TLLM_VALIDATE_CUSTOM_OPS=1`` this returns a plain + ``@torch.library.custom_op`` instead, restoring the per-call schema and + aliasing validation at the cost of the dispatcher tax. """ if "::" not in qualname: raise ValueError(f"qualname must be '::', got {qualname!r}") @@ -74,7 +111,21 @@ def fast_custom_op( 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) + if VALIDATE_CUSTOM_OPS: + return torch.library.custom_op( + qualname, mutates_args=mutates_args_tuple, device_types=device_types + ) + + if device_types is None: + # `custom_op` registers a device-agnostic kernel when device_types is + # omitted; CompositeExplicitAutograd is the low-level equivalent. + dev_types: Tuple[str, ...] = ("CompositeExplicitAutograd",) + else: + names = (device_types,) if isinstance(device_types, str) else tuple(device_types) + # `Library.impl` wants a dispatch key ("CUDA"), while `custom_op` takes + # a device type ("cuda"). Normalize the same way `custom_op` does so + # both spellings work and porting an op cannot break on the casing. + dev_types = tuple(_dispatch_key_for_device(name) for name in names) def decorator(fn: Callable) -> "FastCustomOp": schema = infer_schema(fn, op_name=op_name, mutates_args=mutates_args_tuple) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 8cee91add853..d332b14fb868 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -201,7 +201,7 @@ def forward( ) -@torch.library.custom_op("trtllm::fused_moe", mutates_args=()) +@fast_custom_op("trtllm::fused_moe", mutates_args=(), device_types=None) def fused_moe( input: torch.Tensor, token_selected_experts: torch.Tensor, @@ -510,7 +510,7 @@ def forward( ) -@torch.library.custom_op("trtllm::fp8_rowwise_gemm", mutates_args=()) +@fast_custom_op("trtllm::fp8_rowwise_gemm", mutates_args=(), device_types=None) def fp8_rowwise_gemm( act: torch.Tensor, weight: torch.Tensor, @@ -656,7 +656,9 @@ def forward( ) -@torch.library.custom_op("trtllm::mxfp8_mxfp8_gemm_autotuned", mutates_args=()) +@fast_custom_op("trtllm::mxfp8_mxfp8_gemm_autotuned", + mutates_args=(), + device_types=None) def mxfp8_mxfp8_gemm_autotuned( act: torch.Tensor, act_scale: torch.Tensor, @@ -1012,7 +1014,9 @@ def forward( return result -@torch.library.custom_op("trtllm::nvfp4_gemm_cublaslt", mutates_args=()) +@fast_custom_op("trtllm::nvfp4_gemm_cublaslt", + mutates_args=(), + device_types=None) def nvfp4_gemm_cublaslt( act_fp4: torch.Tensor, weight: torch.Tensor, @@ -1072,7 +1076,9 @@ def _( dtype=output_dtype) -@torch.library.custom_op("trtllm::nvfp4_gemm_cutlass", mutates_args=()) +@fast_custom_op("trtllm::nvfp4_gemm_cutlass", + mutates_args=(), + device_types=None) def nvfp4_gemm_cutlass( act_fp4: torch.Tensor, weight: torch.Tensor, @@ -1603,7 +1609,9 @@ def get_tuning_config(cls, use_deep_seek_fp8: bool, return tuning_config -@torch.library.custom_op("trtllm::fp8_batched_gemm_trtllmgen", mutates_args=()) +@fast_custom_op("trtllm::fp8_batched_gemm_trtllmgen", + mutates_args=(), + device_types=None) def fp8_batched_gemm_trtllmgen( mat1: torch.Tensor, mat2: torch.Tensor, @@ -1673,7 +1681,9 @@ def _( return (fake_out, fake_dq_sfs_c) -@torch.library.custom_op("trtllm::w4a8_mxfp4_fp8_gemm", mutates_args=()) +@fast_custom_op("trtllm::w4a8_mxfp4_fp8_gemm", + mutates_args=(), + device_types=None) def w4a8_mxfp4_fp8_gemm( act_fp8: torch.Tensor, weight: torch.Tensor, @@ -1766,7 +1776,9 @@ def forward( ) -@torch.library.custom_op("trtllm::weight_only_quant_gemm", mutates_args=()) +@fast_custom_op("trtllm::weight_only_quant_gemm", + mutates_args=(), + device_types=None) def weight_only_quant_gemm( activation: torch.Tensor, weight: torch.Tensor, @@ -1865,8 +1877,9 @@ def forward(self, kwargs["bias"], kwargs["zeros"], alpha) -@torch.library.custom_op("trtllm::finegrained_mixed_dtype_gemm", - mutates_args=()) +@fast_custom_op("trtllm::finegrained_mixed_dtype_gemm", + mutates_args=(), + device_types=None) def finegrained_mixed_dtype_gemm( input: torch.Tensor, weight: torch.Tensor, @@ -2050,7 +2063,7 @@ def forward( return output -@torch.library.custom_op("trtllm::fp8_swap_ab_gemm", mutates_args=()) +@fast_custom_op("trtllm::fp8_swap_ab_gemm", mutates_args=(), device_types=None) def fp8_swap_ab_gemm( input: torch.Tensor, weight: torch.Tensor, @@ -2158,7 +2171,9 @@ def get_fp8_block_scaling_gemm_constraint_spec() -> Tuple[ConstraintSpec, ...]: return _get_fp8_block_scaling_gemm_constraint_spec(get_sm_version()) -@torch.library.custom_op("trtllm::fp8_block_scaling_gemm", mutates_args=()) +@fast_custom_op("trtllm::fp8_block_scaling_gemm", + mutates_args=(), + device_types=None) def fp8_block_scaling_gemm( a: torch.Tensor, b: torch.Tensor, @@ -2481,7 +2496,7 @@ def _(input: torch.Tensor, group: List[int]) -> bool: "custom op runtime implementation, not from fake/tracing execution.") -@torch.library.custom_op("trtllm::tunable_allreduce", mutates_args=()) +@fast_custom_op("trtllm::tunable_allreduce", mutates_args=(), device_types=None) def tunable_allreduce( input: torch.Tensor, residual: Optional[torch.Tensor], @@ -2849,7 +2864,9 @@ def _quantize_te(self, return quantized_data, scale -@torch.library.custom_op("trtllm::quantize_e4m3_per_tensor", mutates_args=()) +@fast_custom_op("trtllm::quantize_e4m3_per_tensor", + mutates_args=(), + device_types=None) def quantize_e4m3_per_tensor( input: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """ diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 0063abfa4f1c..6c9604c7ed05 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -28,6 +28,7 @@ from ..autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) +from .fast_custom_op import fast_custom_op _MOE_AUTOTUNE_DUMMY_DISTRIBUTION_ENV = ( "TRTLLM_GEN_MOE_AUTOTUNE_DUMMY_DISTRIBUTION") @@ -647,7 +648,9 @@ def get_tuning_config(cls, return tuning_config -@torch.library.custom_op("trtllm::fp4_block_scale_moe_runner", mutates_args=()) +@fast_custom_op("trtllm::fp4_block_scale_moe_runner", + mutates_args=(), + device_types=None) def fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], routing_bias: Optional[torch.Tensor], hidden_states: torch.Tensor, @@ -1063,7 +1066,9 @@ def get_tuning_config(cls, return tuning_config -@torch.library.custom_op("trtllm::fp8_block_scale_moe_runner", mutates_args=()) +@fast_custom_op("trtllm::fp8_block_scale_moe_runner", + mutates_args=(), + device_types=None) def fp8_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], routing_bias: torch.Tensor, hidden_states: torch.Tensor, @@ -1424,8 +1429,9 @@ def get_tuning_config(cls, return tuning_config -@torch.library.custom_op("trtllm::mxe4m3_mxe2m1_block_scale_moe_runner", - mutates_args=()) +@fast_custom_op("trtllm::mxe4m3_mxe2m1_block_scale_moe_runner", + mutates_args=(), + device_types=None) def mxe4m3_mxe2m1_block_scale_moe_runner( routing_logits: Optional[torch.Tensor], routing_bias: Optional[torch.Tensor], @@ -1745,8 +1751,9 @@ def get_tuning_config(cls, return tuning_config -@torch.library.custom_op("trtllm::e4m3_mxe2m1_block_scale_moe_runner", - mutates_args=()) +@fast_custom_op("trtllm::e4m3_mxe2m1_block_scale_moe_runner", + mutates_args=(), + device_types=None) def e4m3_mxe2m1_block_scale_moe_runner( routing_logits: Optional[torch.Tensor], routing_bias: Optional[torch.Tensor], @@ -2068,8 +2075,9 @@ def get_tuning_config(cls, return tuning_config -@torch.library.custom_op("trtllm::bf16_mxe2m1_block_scale_moe_runner", - mutates_args=()) +@fast_custom_op("trtllm::bf16_mxe2m1_block_scale_moe_runner", + mutates_args=(), + device_types=None) def bf16_mxe2m1_block_scale_moe_runner( routing_logits: Optional[torch.Tensor], routing_bias: Optional[torch.Tensor], @@ -2377,8 +2385,9 @@ def get_tuning_config(cls, return tuning_config -@torch.library.custom_op("trtllm::fp8_fp4_block_scale_moe_runner", - mutates_args=()) +@fast_custom_op("trtllm::fp8_fp4_block_scale_moe_runner", + mutates_args=(), + device_types=None) def fp8_fp4_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], routing_bias: Optional[torch.Tensor], hidden_states: torch.Tensor, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 70d62c0ec386..520fe3a7181d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -26,6 +26,7 @@ l0_cpu: # Kimi K3 disagg parity harness self-test (comparison logic only, no GPUs). - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest - unittest/_torch/attention + - unittest/_torch/custom_ops/test_fast_custom_op.py - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py - unittest/_torch/distributed - unittest/_torch/executor diff --git a/tests/unittest/_torch/custom_ops/test_fast_custom_op.py b/tests/unittest/_torch/custom_ops/test_fast_custom_op.py new file mode 100644 index 000000000000..a8831f66597a --- /dev/null +++ b/tests/unittest/_torch/custom_ops/test_fast_custom_op.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Contract tests for ``fast_custom_op``. + +``fast_custom_op`` trades ``@torch.library.custom_op``'s per-call validation +for a much cheaper dispatch path, so the properties that stop that trade from +becoming a correctness problem are worth pinning down: + +* it registers on the same devices the ported op used to support, +* it accepts both device-type ("cuda") and dispatch-key ("CUDA") spellings, +* it is transparent to ``torch.compile`` (no graph break, same result), and +* ``TLLM_VALIDATE_CUSTOM_OPS=1`` really does restore the aliasing check that + the fast path drops. + +Everything here runs on CPU so the whole file is cheap enough for a +hardware-agnostic CI stage. +""" + +import importlib +import itertools +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.custom_ops import fast_custom_op as fco_module +from tensorrt_llm._torch.custom_ops.fast_custom_op import fast_custom_op + +_COUNTER = itertools.count() + + +def _unique_ns(prefix: str) -> str: + """torch libraries are process-global; give every op a fresh namespace.""" + return f"{prefix}_{next(_COUNTER)}" + + +def test_device_agnostic_registration_runs_on_cpu(): + """``device_types=None`` must keep an op usable on every backend. + + ``@torch.library.custom_op`` registers a backend-agnostic kernel when + ``device_types`` is omitted. Porting such an op must not silently narrow it + to CUDA, which would turn a CPU call into a NotImplementedError. + """ + ns = _unique_ns("fco_agnostic") + + @fast_custom_op(f"{ns}::add_one", mutates_args=(), device_types=None) + def add_one(x: torch.Tensor) -> torch.Tensor: + return x + 1 + + @add_one.register_fake + def _(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + x = torch.arange(4, dtype=torch.float32) + torch.testing.assert_close(getattr(torch.ops, ns).add_one(x), x + 1) + + +@pytest.mark.parametrize("device_types", ["cuda", "CUDA"]) +def test_accepts_device_type_and_dispatch_key_spellings(device_types): + """Both "cuda" and "CUDA" must register. + + ``custom_op`` takes a device type, ``Library.impl`` takes a dispatch key. + Ops in the tree use both spellings, so rejecting either would break a port + at import time. + """ + ns = _unique_ns("fco_spelling") + + @fast_custom_op(f"{ns}::noop", mutates_args=(), device_types=device_types) + def noop(x: torch.Tensor) -> torch.Tensor: + return x.clone() + + @noop.register_fake + def _(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + # Registration alone is the assertion; the op resolves off torch.ops. + assert hasattr(getattr(torch.ops, ns), "noop") + + +def test_matches_custom_op_under_eager_and_compile(): + """Same function, both decorators, same answers -- and no graph break.""" + ns_fast, ns_ref = _unique_ns("fco_fast"), _unique_ns("fco_ref") + + def impl(x: torch.Tensor, scale: float, bias: Optional[torch.Tensor]) -> torch.Tensor: + out = x * scale + return out if bias is None else out + bias + + fast_op = fast_custom_op(f"{ns_fast}::scale", mutates_args=(), device_types=None)(impl) + fast_op.register_fake(lambda x, scale, bias: torch.empty_like(x)) + + ref_op = torch.library.custom_op(f"{ns_ref}::scale", mutates_args=())(impl) + ref_op.register_fake(lambda x, scale, bias: torch.empty_like(x)) + + x = torch.randn(8) + bias = torch.randn(8) + fast_call = lambda: getattr(torch.ops, ns_fast).scale(x, 2.0, bias) # noqa: E731 + ref_call = lambda: getattr(torch.ops, ns_ref).scale(x, 2.0, bias) # noqa: E731 + + torch.testing.assert_close(fast_call(), ref_call()) + + explained = torch._dynamo.explain(fast_call)() + assert explained.graph_break_count == 0 + + torch._dynamo.reset() + torch.testing.assert_close(torch.compile(fast_call, dynamic=False)(), ref_call()) + + +def test_validate_env_restores_aliasing_check(monkeypatch): + """``TLLM_VALIDATE_CUSTOM_OPS=1`` must catch an undeclared alias. + + An op that declares ``mutates_args=()`` but returns an input aliases that + input. The fast path cannot see this and would let a later in-place write + corrupt the caller's tensor; the validating path must raise instead. This + is the whole reason the switch exists, so assert both halves. + """ + ns_fast = _unique_ns("fco_alias_fast") + + @fast_custom_op(f"{ns_fast}::identity", mutates_args=(), device_types=None) + def identity(x: torch.Tensor) -> torch.Tensor: + return x # undeclared alias of the input + + @identity.register_fake + def _(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + # Fast path: no check, so the aliasing write lands on the caller's tensor. + x = torch.zeros(3) + getattr(torch.ops, ns_fast).identity(x).add_(1.0) + assert torch.equal(x, torch.ones(3)), "fast path is expected to alias" + + # Validating path: same bug, now rejected. + monkeypatch.setenv("TLLM_VALIDATE_CUSTOM_OPS", "1") + validating = importlib.reload(fco_module) + assert validating.VALIDATE_CUSTOM_OPS + + ns_val = _unique_ns("fco_alias_val") + + @validating.fast_custom_op(f"{ns_val}::identity", mutates_args=(), device_types=None) + def identity_validated(x: torch.Tensor) -> torch.Tensor: + return x + + @identity_validated.register_fake + def _(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + with pytest.raises(RuntimeError, match="alias"): + getattr(torch.ops, ns_val).identity(torch.zeros(3)) + + # Leave the module in its default (fast) state for other tests. + monkeypatch.delenv("TLLM_VALIDATE_CUSTOM_OPS") + importlib.reload(fco_module) From 95f9ba3a2b089b723283728251009829c2f5eb76 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Fri, 4 Sep 2026 02:09:43 +0000 Subject: [PATCH 2/2] [None][fix] Mark test_fast_custom_op as cpu_only CI runs `l0_cpu.yml` entries with `-m cpu_only`, and `tests/unittest/conftest.py` additionally ignores any collected file whose source lacks the literal `pytest.mark.cpu_only`. The new file had neither, so all 5 tests were deselected and pytest exited 5, which the unittest wrapper reports as a failure. This hit both `CPU-Generic-x86-1` and `CPU-Generic-arm-1` in PR_Github #71213. Verified: `pytest -m cpu_only` on the file now selects and passes all 5 tests where it previously reported `5 deselected / 0 selected`. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tests/unittest/_torch/custom_ops/test_fast_custom_op.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/_torch/custom_ops/test_fast_custom_op.py b/tests/unittest/_torch/custom_ops/test_fast_custom_op.py index a8831f66597a..f8d66a8f7211 100644 --- a/tests/unittest/_torch/custom_ops/test_fast_custom_op.py +++ b/tests/unittest/_torch/custom_ops/test_fast_custom_op.py @@ -26,6 +26,10 @@ from tensorrt_llm._torch.custom_ops import fast_custom_op as fco_module from tensorrt_llm._torch.custom_ops.fast_custom_op import fast_custom_op +# Registered in l0_cpu.yml, which CI runs with `-m cpu_only`; without this the +# whole file is deselected and pytest exits 5. +pytestmark = pytest.mark.cpu_only + _COUNTER = itertools.count()