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..f8d66a8f7211 --- /dev/null +++ b/tests/unittest/_torch/custom_ops/test_fast_custom_op.py @@ -0,0 +1,155 @@ +# 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 + +# 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() + + +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)