[TRTLLM-11958][perf] reduce @torch.library.custom_op host overhead by luyiyun1021 · Pull Request #13149 · NVIDIA/TensorRT-LLM

luyiyun1021 · 2026-04-17T09:19:27Z

Description

@torch.library.custom_op wraps every call in Python-level wrappers registered on several dispatch keys (Autograd, optionally ADInplaceOrView, and the backend key), which imposes a ~7us per-call dispatcher tax on top of the actual kernel launch. On host-heavy iterations (LTX-2 dense transformer issues ~2100 such calls per step for the two ops in this PR) the tax becomes a measurable fraction of the per-step wall time.

This PR:

Introduces a thin helper fast_custom_op that registers an op directly through the low-level torch.library.Library.define + impl API. The helper preserves the @custom_op developer experience — schema is inferred from Python type hints via torch.library.infer_schema, and the returned object exposes .register_fake — while bypassing the multi-layer Python wrappers that @custom_op installs.
Migrates two hot ops on the LTX-2 VisualGen path (trtllm::nvfp4_gemm and trtllm::tunable_fp4_quantize) to @fast_custom_op.

Approach — usage

# Before
@torch.library.custom_op("trtllm::nvfp4_gemm", mutates_args=())
def nvfp4_gemm(act_fp4: torch.Tensor, ...) -> torch.Tensor: ...
@nvfp4_gemm.register_fake
def _(...): ...

# After (same ergonomics, ~5us/call cheaper)
from tensorrt_llm._torch.custom_ops.fast_custom_op import fast_custom_op

@fast_custom_op("trtllm::nvfp4_gemm", mutates_args=())
def nvfp4_gemm(act_fp4: torch.Tensor, ...) -> torch.Tensor: ...
@nvfp4_gemm.register_fake
def _(...): ...

The helper uses torch.library.infer_schema internally so the schema is still driven by Python type hints — no hand-written schema strings. FRAGMENT mode is used under the hood because the trtllm namespace is already declared by C++ via TORCH_LIBRARY_FRAGMENT.

Why `@custom_op` is expensive — code-level analysis

Looking at torch/_library/custom_ops.py::CustomOpDef._register_to_dispatcher (L607-675), @custom_op registers several Python-level kernels on multiple dispatch keys. Each one is executed on every call:

1. Autograd key (always registered) — torch/_library/autograd.py::autograd_impl (L108):

def autograd_impl(keyset, *args, **keyword_only_args):
    if _C.is_grad_enabled() and _C._any_requires_grad(*args):   # tensor iteration
        result = Generated.apply(*args, Metadata(keyset, kwargs))
    else:
        result = forward_no_grad(*args, Metadata(keyset, kwargs))
    return result

def forward_no_grad(*args):
    metadata = args[-1]; args = args[:-1]
    with _C._AutoDispatchBelowAutograd():                       # Python ctx mgr
        return op.redispatch(                                    # re-dispatch
            keyset & _C._after_autograd_keyset, *args, **metadata.keyword_only_args)

Every call pays: is_grad_enabled() + _any_requires_grad(*args) tensor iteration + Metadata dataclass construction + _AutoDispatchBelowAutograd context manager + op.redispatch(...) (a second dispatch trip).

2. ADInplaceOrView key — adinplaceorview_impl (L654). Registered only when the schema is mutable or a view op. Bumps version counters on mutated args and routes through call_boxed. Not a cost for the two ops in this PR (both are pure, mutates_args=()), so schema is non-mutable and this wrapper is not installed.

3. CUDA backend key — backend_impl wrapper from register_kernel (L346-362):

def backend_impl(*args, **kwargs):
    result = self._backend_fns[device_type](*args, **kwargs)    # user's fn
    def get_module():                                            # closure
        return inspect.getmodule(self._backend_fns[device_type])
    schema = self._opoverload._schema
    if not schema._is_view_op():
        utils._c_check_aliasing_constraint(                       # aliasing check
            self._name, args, kwargs, result, get_module)
    return result

Every call pays an aliasing-constraint check against self._opoverload._schema (iterates inputs/outputs, compares storage pointers) and a closure construction.

4. CustomOpDef.__call__ — L697. One extra Python frame when the op is invoked by the decorated name (e.g. nvfp4_gemm(x) in-module). Call sites that go through torch.ops.trtllm.nvfp4_gemm(x) bypass this frame.

What the low-level `Library.define + impl` path does

lib = Library("trtllm", "FRAGMENT")
lib.define(schema_str)       # declare op
lib.impl(name, fn, "CUDA")   # user's fn registered directly on the CUDA key

No Autograd kernel registered → a call with requires_grad=True falls through to the dispatcher's C++-level "no autograd kernel" path (same user-facing error as before), no Python wrapper runs on the common inference path where no tensor needs grad.
No ADInplaceOrView wrapper.
CUDA key points at the user function directly; no backend_impl wrapper, no aliasing check on each call.

Call path comparison (same `torch.ops.trtllm.my_op(x)` invocation)

@custom_op path (pure, non-mutating op):

C++ dispatcher
 └─ Autograd key
     └─ autograd_impl(keyset, x)                        [Python]
         ├─ is_grad_enabled(); _any_requires_grad(*args)
         ├─ Metadata(keyset, kwargs)
         └─ forward_no_grad(x, Metadata)                [Python]
              └─ with _AutoDispatchBelowAutograd():     [Python ctx mgr]
                  └─ op.redispatch(..., x)              [re-dispatch]
                       └─ C++ dispatcher: CUDA key
                            └─ backend_impl(x)          [Python]
                                 ├─ user_fn(x)
                                 ├─ inspect.getmodule(...) closure
                                 └─ _c_check_aliasing_constraint(...)

Library.define + impl path:

C++ dispatcher
 └─ CUDA key
     └─ user_fn(x)

Two fewer Python frames + no aliasing check + no re-dispatch.

Microbenchmark

Isolated per-call cost on B200 + PyTorch 2.10, 20k-iter tight loop with x.clone() as a minimal kernel (tmp/bench_custom_op_overhead_v2.py):

Registration	Total per-call	Pure dispatcher tax*	vs `@custom_op`
plain Python fn (kernel floor)	4.72us	—	—
`@torch.library.custom_op` (baseline)	11.67us	7.02us	—
Manual `Library.define + impl`	6.28us	1.63us	−5.39us
`@fast_custom_op` (via `torch.ops`)	6.28us	1.63us	−5.39us ✅
`@fast_custom_op` (via proxy `__call__`)	6.09us	1.44us	−5.57us

*Pure dispatcher tax = total per-call minus the plain-Python-fn kernel floor.

Key observations:

@fast_custom_op's hot path (torch.ops.trtllm.<name>(...)) is byte-identical to manual Library.define + impl at runtime: both resolve to the same OpOverload and skip all the Python wrappers. No wrapper overhead from the helper itself.
The proxy __call__ path is marginally cheaper because the OpOverload object is cached at decoration time.

Estimated contribution on LTX-2 at baseline (@custom_op, 12us/call gross):

Op	Calls per step	Saving per step
`trtllm::tunable_fp4_quantize`	~1260	~5.4us × 1260 ≈ 6.8 ms
`trtllm::nvfp4_gemm`	~840	~5.4us × 840 ≈ 4.5 ms
Combined	~2100	~11.3 ms / step

End-to-end measurements

Config	Before	After	Delta
1 GPU, 40-step, 736×1280, 121f	33.95 s	33.61 s	−1.0 %
8 GPU (cfg=2, uly=4), 40-step	10.21 s	9.67 s	−5.3 %
nsys Pure CPU idle / step	baseline	−29 ms/step	—

Per-step host-time attribution, 1 GPU non-cuda-graph vs compile modes (from nsys, denoise_step NVTX range, baseline @custom_op state):

Config	Per-step host active	2 ops' `@custom_op` contribution (12us × 2100 ≈ 25.2 ms)
Eager, no compile, no graph	~1007 ms	~2.5 %
torch.compile, no cuda_graph	~548 ms	~4.6 %
torch.compile + cuda_graph	~193 ms (from 8 GPU trace)	~13 %

The E2E win scales with how much other host overhead (kernel launch APIs) has already been absorbed by torch.compile + CUDA Graph — on torch.compile + cuda_graph the @custom_op wrapper becomes a large fraction of what's left.

When to keep `@custom_op`

Feature	Needed when...	Used by these ops?
`mutates_args=("out",)`	Kernel writes in-place to an input buffer (e.g. `trtllm::bmm_out`)	No — both ops return fresh tensors
`setup_context` + `backward`	Op is part of a training graph	No — inference-only
auto-functionalization	Rewrites mutating calls inside `torch.compile` — only fires when `mutates_args` is non-empty	No — `mutates_args=()` makes this a no-op

Keep @torch.library.custom_op for ops that (a) have non-empty mutates_args, (b) need autograd, or (c) are under active development and benefit from the richer Python-side error messages.

Functional equivalence

After the switch, torch.ops.trtllm.nvfp4_gemm and torch.ops.trtllm.tunable_fp4_quantize resolve to the same OpOverload objects with the same schema string as before — so eager Python, torch.compile, Dynamo, FX passes, and any downstream FX pattern matcher (e.g. ar_residual_norm that looks for torch.ops.trtllm.nvfp4_gemm.default) all see an identical op. The only difference is the dispatch path, which is shorter. register_fake still provides the same FakeTensor/meta shape inference used by torch.compile tracing.

Side-effect validation

Verified parity with @custom_op baseline via 13 targeted tests: numerical correctness, torch.compile (fullgraph + dynamic), FakeTensor/meta propagation, inference_mode / no_grad, mutation-safety (args not mutated, output not aliased to inputs), error cases (wrong dtype/device) raising identical RuntimeError, CPU-tensor dispatch failing identically, autograd-path with requires_grad=True raising the same error (both ops are non-differentiable by design).

E2E smoke test: 10-step LTX-2 1 GPU run with torch_compile=true using the @fast_custom_op form — exit code 0, steady-state 0.75s/step matches the manual-Library form.

Test Coverage

tests/unittest/_torch/thop/test_nvfp4_gemm.py covers nvfp4_gemm — passes unchanged.
tunable_fp4_quantize is exercised by all NVFP4 Linear / MoE tests and the LTX-2 integration.
No new test surface: registration-API change only; behavior and schema are preserved.

PR Checklist

PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

luyiyun1021 · 2026-04-17T09:19:31Z

/bot run --disable-fail-fast

coderabbitai · 2026-04-17T09:22:58Z

📝 Walkthrough

Walkthrough

Refactored custom torch operation registration for nvfp4_gemm and tunable_fp4_quantize from decorator-based @torch.library.custom_op to explicit torch.library.Library() API calls with separate schema definitions, CUDA implementations, and fake function registrations.

Changes

Cohort / File(s)	Summary
Custom Torch Op Registration `tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`	Migrated two custom operations from `@torch.library.custom_op` decorators to explicit Library API registration using `.define()`, `.impl(..., "CUDA")`, and `torch.library.register_fake()` calls; renamed internal fake functions (`_nvfp4_gemm_fake`, `_tunable_fp4_quantize_fake`) for clarity.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description is comprehensive, providing detailed technical context, code analysis, performance metrics, and clear before/after comparisons with well-documented trade-offs.
Title check	✅ Passed	The title clearly and specifically describes the main change: refactoring custom operator registration to reduce decorator overhead for performance.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
2367-2374: Schema lacks default values present in Python function.

The Python function tunable_fp4_quantize has default values (scaling_vector_size=16, is_sf_swizzled_layout=False), but the schema definition doesn't include them. This means callers using torch.ops.trtllm.tunable_fp4_quantize must always provide all four arguments.

From the relevant code snippet in linear.py, all arguments are passed explicitly, so this works. However, for API consistency with the Python function and other ops in this file, consider adding defaults to the schema:
 _trtllm_tunable_fp4_quantize_lib.define(
-    "tunable_fp4_quantize(Tensor input, Tensor input_scale, "
-    "int scaling_vector_size, bool is_sf_swizzled_layout) -> Tensor[]")
+    "tunable_fp4_quantize(Tensor input, Tensor input_scale, "
+    "int scaling_vector_size=16, bool is_sf_swizzled_layout=False) -> Tensor[]")
If all call sites explicitly pass these arguments, this is a minor inconsistency rather than a functional issue.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py` around lines 2367 - 2374,
The schema registered via _trtllm_tunable_fp4_quantize_lib.define for the op
tunable_fp4_quantize currently omits the Python defaults
(scaling_vector_size=16, is_sf_swizzled_layout=False); update the schema string
passed to _trtllm_tunable_fp4_quantize_lib.define to include those default
values so torch.ops.trtllm.tunable_fp4_quantize supports the same optional
arguments as the Python function (i.e., add "=16" to scaling_vector_size and
"=False" to is_sf_swizzled_layout in the define call).

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 2367-2374: The schema registered via
_trtllm_tunable_fp4_quantize_lib.define for the op tunable_fp4_quantize
currently omits the Python defaults (scaling_vector_size=16,
is_sf_swizzled_layout=False); update the schema string passed to
_trtllm_tunable_fp4_quantize_lib.define to include those default values so
torch.ops.trtllm.tunable_fp4_quantize supports the same optional arguments as
the Python function (i.e., add "=16" to scaling_vector_size and "=False" to
is_sf_swizzled_layout in the define call).

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 0abef18f-78dc-4cb9-b5cf-e0cc434bc0f6

📥 Commits

Reviewing files that changed from the base of the PR and between 2a0bcb1 and fc8cabb.

📒 Files selected for processing (1)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

tensorrt-cicd · 2026-04-17T09:25:18Z

PR_Github #44012 [ run ] triggered by Bot. Commit: fc8cabb Link to invocation

tensorrt-cicd · 2026-04-18T09:26:07Z

PR_Github #44012 [ run ] completed with state ABORTED. Commit: fc8cabb

Link to invocation

hyukn

This is a good idea for surgically eliminating the overhead caused by torch.custom_op wrapper system. But it sacrifices the original Python-level friendly developer guardrails provided by @custom_op (like automatic type inference), and it may require more constraints for the op author.

Another way is to apply a well-defined rule periodically for the existing ops, "translate" them into the way you did for these two ops, and guarantee the correctness. This may be safer.

luyiyun1021 · 2026-04-22T04:47:26Z

This is a good idea for surgically eliminating the overhead caused by torch.custom_op wrapper system. But it sacrifices the original Python-level friendly developer guardrails provided by @custom_op (like automatic type inference), and it may require more constraints for the op author.

Another way is to apply a well-defined rule periodically for the existing ops, "translate" them into the way you did for these two ops, and guarantee the correctness. This may be safer.

Thanks. @hyukn Could you please take a look at my newest commit to see if it addressed your concern?

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>

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.<name>(...) goes through the C++ dispatcher directly, same as the manual form). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

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>

luyiyun1021 · 2026-04-22T11:23:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-22T12:00:04Z

PR_Github #44962 [ run ] triggered by Bot. Commit: bc09821 Link to invocation

luyiyun1021 · 2026-04-23T01:44:58Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-23T01:50:48Z

PR_Github #45063 [ run ] triggered by Bot. Commit: bc09821 Link to invocation

hyukn

LGTM. This is more modular-designed. Let us see if CI reports any potential issues.

luyiyun1021 · 2026-04-23T15:41:02Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-23T15:46:52Z

PR_Github #45210 [ run ] triggered by Bot. Commit: bc09821 Link to invocation

tensorrt-cicd · 2026-04-23T23:22:59Z

PR_Github #45210 [ run ] completed with state SUCCESS. Commit: bc09821
/LLM/main/L0_MergeRequest_PR pipeline #35477 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

luyiyun1021 · 2026-04-24T02:19:59Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-24T02:26:58Z

PR_Github #45299 [ run ] triggered by Bot. Commit: bc09821 Link to invocation

tensorrt-cicd · 2026-04-24T06:19:45Z

PR_Github #45299 [ run ] completed with state SUCCESS. Commit: bc09821
/LLM/main/L0_MergeRequest_PR pipeline #35553 completed with status: 'SUCCESS'

CI Report

Link to invocation

…VIDIA#13149) Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

luyiyun1021 requested a review from a team as a code owner April 17, 2026 09:19

luyiyun1021 requested a review from liji-nv April 17, 2026 09:19

github-actions Bot assigned luyiyun1021 Apr 17, 2026

coderabbitai Bot reviewed Apr 17, 2026

View reviewed changes

luyiyun1021 marked this pull request as draft April 17, 2026 09:29

hyukn reviewed Apr 20, 2026

View reviewed changes

luyiyun1021 marked this pull request as ready for review April 22, 2026 04:45

luyiyun1021 added 3 commits April 22, 2026 11:22

luyiyun1021 force-pushed the reduce-custom-op-host-overhead branch from dc3ca46 to bc09821 Compare April 22, 2026 11:23

hyukn approved these changes Apr 23, 2026

View reviewed changes

hyukn mentioned this pull request Apr 23, 2026

[TRTLLM-10288][perf] Reduce AutoTuner host overhead in inference hot path #13116

Closed

1 task

luyiyun1021 changed the title ~~[None][perf] reduce @torch.library.custom_op host overhead~~ [TRTLLM-11958][perf] reduce @torch.library.custom_op host overhead Apr 24, 2026

luyiyun1021 merged commit 95db868 into NVIDIA:main Apr 24, 2026
7 checks passed

yufeiwu-nv pushed a commit to yufeiwu-nv/TensorRT-LLM that referenced this pull request May 19, 2026

[TRTLLM-11958][perf] reduce @torch.library.custom_op host overhead (N…

54fa21c

…VIDIA#13149) Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

coderabbitai Bot mentioned this pull request Jun 15, 2026

Add docs folder NVIDIA/cudnn-frontend#308

Merged

hyukn mentioned this pull request Sep 3, 2026

[None][perf] Port tunable custom ops to fast_custom_op #18645

Open

1 task

Uh oh!

Conversation

luyiyun1021 commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Approach — usage

Why @custom_op is expensive — code-level analysis

What the low-level Library.define + impl path does

Call path comparison (same torch.ops.trtllm.my_op(x) invocation)

Microbenchmark

End-to-end measurements

When to keep @custom_op

Functional equivalence

Side-effect validation

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

luyiyun1021 commented Apr 17, 2026

Uh oh!

coderabbitai Bot commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Apr 17, 2026

Uh oh!

tensorrt-cicd commented Apr 18, 2026

Uh oh!

hyukn left a comment

Choose a reason for hiding this comment

Uh oh!

luyiyun1021 commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

luyiyun1021 commented Apr 22, 2026

Uh oh!

tensorrt-cicd commented Apr 22, 2026

Uh oh!

luyiyun1021 commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

hyukn left a comment

Choose a reason for hiding this comment

Uh oh!

luyiyun1021 commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

tensorrt-cicd commented Apr 23, 2026

Uh oh!

luyiyun1021 commented Apr 24, 2026

Uh oh!

tensorrt-cicd commented Apr 24, 2026

Uh oh!

tensorrt-cicd commented Apr 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

luyiyun1021 commented Apr 17, 2026 •

edited

Loading

Why `@custom_op` is expensive — code-level analysis

What the low-level `Library.define + impl` path does

Call path comparison (same `torch.ops.trtllm.my_op(x)` invocation)

When to keep `@custom_op`

coderabbitai Bot commented Apr 17, 2026 •

edited

Loading

luyiyun1021 commented Apr 22, 2026 •

edited

Loading