[None][fix] Fix FP8 per-tensor torch.compile graph break in dynamic quantization

[karljang]

Summary

The tensorrt_llm::quantize_e4m3_per_tensor lacks a register_fake implementation. Without register_fake, torch.compile's Dynamo tracer cannot infer output shape/dtype metadata, causing a graph break at every dynamic quantization call.

Added register_fake for tensorrt_llm::quantize_e4m3_per_tensor in cpp_custom_ops.py, matching the pattern already used for the static variant (static_quantize_e4m3_per_tensor).

Test plan

• Verify FP8 per-tensor eager mode produces identical outputs (same kernel)
• Verify torch.compile produces single monolithic FX graph
• Run existing FP8 per-tensor unit tests

Observation

Impact on FLUX.2 (B200, 1024x1024, 50 steps, torch.compile):

• Before: 36 subgraphs, 491 traced nodes (~3% compile coverage)
• After: 1 subgraph, 6,431 traced nodes (full compile coverage)

No latency change observed (GEMMs dominate runtime), but the fix produces a correct monolithic FX graph that enables future Inductor optimizations requiring whole-graph visibility.

Summary by CodeRabbit

• Chores
  • Internal update to quantization operation namespace path. No changes to user-facing functionality or behavior.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

A namespace path is updated for a per-tensor quantization operator in the linear module. The call changes from torch.ops.tensorrt_llm.quantize_e4m3_per_tensor to torch.ops.trtllm.quantize_e4m3_per_tensor, with arguments and return values remaining unchanged.

Changes

Cohort / File(s)	Summary
Quantization Operator Namespace Update tensorrt_llm/_torch/modules/linear.py	Updated operator namespace from torch.ops.tensorrt_llm to torch.ops.trtllm for the quantize_e4m3_per_tensor call during dynamic quantization.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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
Title check	✅ Passed	Title clearly describes the main change: fixing FP8 per-tensor torch.compile graph fragmentation in dynamic quantization by switching operator namespaces.
Description check	✅ Passed	PR description clearly explains the issue, solution, and test coverage with specific metrics demonstrating the fix's impact.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

[karljang]

/bot run

[tensorrt-cicd]

PR_Github #36981 [ run ] triggered by Bot. Commit: 2bdd317 Link to invocation

[tensorrt-cicd]

PR_Github #36981 [ run ] completed with state SUCCESS. Commit: 2bdd317
/LLM/main/L0_MergeRequest_PR pipeline #28634 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

Link to invocation

[karljang]

Merged as discussed with @liji-nv ~ thank you!