Enable torch.compile under a flag by demandal25 · Pull Request #210 · AMD-Ecosystem/flashinfer

demandal25 · 2026-03-31T14:40:40Z

Adds an environment flag FLASHINFER_USE_TORCH_CUSTOM_OPS, setting which enables the torch.compile in flashinfer. By default, it's disabled in the upstream (https://github.com/ROCm/flashinfer/blob/61a9b747f5ad0c9a886879b4d8e4d1ef046dcf39/flashinfer/utils.py#L273).

How to use it

`FLASHINFER_USE_TORCH_CUSTOM_OPS=1 <your command>

Testing

Adds a pytest for it. tests/rocm_tests/test_torch_compile_hip.py to be run automatically when we run pytest.

All pytests passed for both the cases:

FLASHINFER_USE_TORCH_CUSTOM_OPS=0 pytest, or simply pytest, and
FLASHINFER_USE_TORCH_CUSTOM_OPS=1 pytest

Copilot

Pull request overview

Adds opt-in support for torch.compile by wrapping FlashInfer kernels in opaque torch.library.custom_ops (to avoid Dynamo tracing into extension code), plus a small ROCm benchmark/verification script and optional local Git hook scaffolding.

Changes:

Introduce FLASHINFER_USE_TORCH_CUSTOM_OPS env flag (checked at import time) to enable/disable torch.library.custom_op wrapping.
Export use_torch_custom_ops_enabled() from the top-level flashinfer package for both CUDA and ROCm/HIP builds.
Add a ROCm micro-benchmark script and optional .githooks pre-push protection for amd-integration.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`flashinfer/utils.py`	Adds env-gated custom-op enablement + helper to query the setting; conditionally registers `custom_op`/`register_fake`.
`flashinfer/__init__.py`	Re-exports `use_torch_custom_ops_enabled()` for CUDA and HIP branches.
`scripts/verify_enable_torch_compile.py`	New ROCm-focused eager vs `torch.compile` micro-benchmark for `append_paged_kv_cache`.
`.githooks/README`	Documents how to enable optional local Git hooks and intended behavior.
`.githooks/pre-push`	New local pre-push hook intended to block pushing/updating `amd-integration`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Made-with: Cursor

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…ytest - Remove contextlib.suppress in register_custom_op so registration failures surface instead of silently falling back (Copilot + diptorupd) - Fold torch >= 2.4 check into _USE_TORCH_CUSTOM_OPS so use_torch_custom_ops_enabled() reports effective behavior (Copilot) - Replace scripts/verify_enable_torch_compile.py with tests/rocm_tests/test_torch_compile_hip.py covering eager and torch.compile paths in subprocess isolation (diptorupd) Made-with: Cursor

Drop the repo-root include fallback per reviewer feedback: running from a source tree without an install is not a supported workflow. Made-with: Cursor

- Use device="cuda" (not "hip") — PyTorch on ROCm uses "cuda" as the device string - Check HIP availability properly in pytestmark - Accept any exception from torch.compile when custom ops are disabled, since the actual error (TorchRuntimeError from FakeTensor) differs from the _guard_compile message Made-with: Cursor

torch.library.custom_op cannot infer schemas for parameters like Optional[torch.Generator]. Instead of crashing, fall back to _guard_compile so these ops still get torch.compile protection. Made-with: Cursor

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

diptorupd

Thank you! LGTM.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

…ag (#230) ## Summary Follow-up to #210 ("Enable torch.compile under a flag"), addressing Copilot and reviewer comments. - **`plan_info_vec_as_tensor` always returns CPU tensor**: The ROCm C++ ops read `plan_info_vec.data_ptr<int64_t>()` from the host side, so the tensor must stay on CPU regardless of where the workspace buffer lives. The old code moved it to the GPU device, which is semantically wrong even if it happened to work on MI300X's unified memory. The `device` parameter is kept for API compatibility but now ignored. - **`register_custom_op` warns on schema fallback**: Added `warnings.warn` in the `ValueError`/`TypeError` catch so the fallback to `_guard_compile` is visible. The fallback behavior is kept — `torch.compile` still raises a clear `RuntimeError` if it traces into an op that couldn't be registered as a custom op. - **`test_eager_with_custom_ops` gets `skipif(torch<2.4)`**: `_USE_TORCH_CUSTOM_OPS` is already gated on `torch >= 2.4`, so asserting it's `True` would fail on older torch. Adds the same skip condition the other compile tests use. - **Tensor-core decode path gets a None guard**: Added `assert self._plan_info is not None` before building `run_args` in the tensor-core branch of `BatchDecodeWithPagedKVCacheWrapper.forward()`, matching the non-tensor-core branch. ## Test plan - [x] `pytest tests/rocm_tests/test_torch_compile_hip.py -v` — 4 passed - [x] `pytest tests/rocm_tests/ -m "not slow"` — full suite in progress --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

## Summary Refresh the FlashInfer+ROCm README aimed at library consumers, refresh the Feature Support Matrix to match what has actually landed on `amd-integration`, and align the ROCm MLA wrapper with the rest of the ROCm backends so `backend="auto"` is accepted everywhere. ### What changed #### `README.md` - **Intro and structure.** Tighten the intro to call out HIP-in-repo kernels vs AITER dispatch up front; link to the Feature Support Matrix and AITER sections from the first paragraph. Cross-link CDNA3 / CDNA4 to AMD's official architecture whitepapers on first mention. - **Feature Support Matrix.** Replaced with a five-column table (Kernel / HIP / AITER / `backend="auto"` resolves to / Notes). New ✅ rows: Cascade (#221), MLA via AITER (#232), RoPE (#223), paged KV-cache append, RMSNorm via AITER (#232), sliding-window decode on the AITER path (#234), activation, quantization, and opt-in `torch.compile` (#210). Every ✅ is backed by a `tests/rocm_tests/test_*_hip.py`. FP8 status is folded into per-row notes rather than a dedicated column. - **GPU / ROCm / PyTorch.** Consolidated into one section with arch codenames inline (gfx942 → MI300X/MI325X = CDNA3, gfx950 → MI355X = CDNA4). `pip install torch` uses `--index-url` instead of `-f` so pip cannot silently fall back to a CPU-only PyPI wheel (matches CLAUDE.md). - **Getting Started.** Collapsed the Docker image table to the latest validated tag and pointed at Docker Hub for older releases. Dropped the manual `micromamba activate base` step (the env is auto-activated). Used the concrete image tag plus a `--name=flashinfer-rocm` in the `docker run` snippet. - **Trying the Examples.** Simplified to point at `examples/` plus one run command — no wget-based downloads. - **Install from Source.** Renamed from "Build from Source"; rewrote the ambiguous "Environment name varies …" note (and later removed it once the build / run blocks made the matching tag self-evident). - **AITER Support.** Collapsed the section intro to avoid re-listing conditions already in the matrix; cross-link Known Limitations. Rewrote Known Limitations preamble to state the two-group split (hard errors vs silently-ignored kwargs). Dropped the redundant Single Prefill Example (Basic Usage already shows the call pattern). - **Environment Variables.** New section documenting runtime env vars — `FLASHINFER_USE_TORCH_CUSTOM_OPS`, `FLASHINFER_HIP_FUSED_CASCADE`, `FLASHINFER_LOGGING_LEVEL`, `FLASHINFER_DISABLE_JIT`, `ROCM_PATH` / `ROCM_HOME`. Build-time vars stay in `CLAUDE.md` and are linked from here. - **Runtime Helpers.** Short snippet showing `is_aiter_supported` and `check_torch_rocm_compatibility`; calls out `validate_flashinfer_rocm_arch` as a build-time validator, not a runtime helper. - **CPX-mode pytest notes.** Split the dense paragraph into labelled bullets (Worker count / Reruns / `slow` marker / HIPBLAS retry). - **Basic Usage.** Moved to the end of the README as a closing example. - **License and Acknowledgements.** Added; the contributing reminder lives on its own line. #### `flashinfer/mla_rocm.py` + `tests/rocm_tests/test_mla_aiter_hip.py` - Accept `backend="auto"` as an alias for `"aiter"` on the ROCm MLA wrapper (default is now `"auto"` to match every other ROCm wrapper). Previously the wrapper raised `ValueError` on anything other than `"aiter"`, leaving MLA as the odd one out in the public API even though there is exactly one implementation to pick from on ROCm. - New tests: `test_mla_backend_accepts_auto_and_aiter` (parametrized over both values) and `test_mla_backend_rejects_unsupported` (confirms `backend="fa2"` still raises; runs without a GPU since the check fires before the AITER probe). ## Test plan - [x] `pre-commit run -a` passes. - [x] `pre-commit run markdownlint --files README.md` passes after every change. - [x] Every TOC entry resolves to an `##` heading in the body. - [x] Every ✅ in the Feature Support Matrix has a backing `tests/rocm_tests/test_*_hip.py`. - [x] `pytest tests/rocm_tests/test_mla_aiter_hip.py` — 11 passed. - [x] Render the README on the PR page and visually confirm tables, code blocks, and `<details>` sections look right. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings March 31, 2026 14:40

demandal25 force-pushed the enable-torch-compile branch from 4518269 to 8f2cd78 Compare March 31, 2026 14:45

Copilot AI reviewed Mar 31, 2026

View reviewed changes

Comment thread .githooks/pre-push Outdated

Comment thread flashinfer/utils.py Outdated

Comment thread scripts/verify_enable_torch_compile.py Outdated

demandal25 force-pushed the enable-torch-compile branch from 8f2cd78 to 768e946 Compare March 31, 2026 14:54

Enable torch.compile under a flag

4e3bb90

demandal25 force-pushed the enable-torch-compile branch from 768e946 to 4e3bb90 Compare March 31, 2026 15:00

demandal25 marked this pull request as draft March 31, 2026 15:01

demandal25 added 4 commits April 2, 2026 00:40

Fix to enable torch.compile

b1bc427

Revert cuda related changes

f4da9a9

Refactor the changelist

4becdb0

Made-with: Cursor

Fix sampling

c9ef604

demandal25 force-pushed the enable-torch-compile branch from 9eadd10 to c9ef604 Compare April 2, 2026 01:31

error handling when torch.compile called witout env

c7b0ef2

demandal25 marked this pull request as ready for review April 3, 2026 02:48

Copilot AI review requested due to automatic review settings April 3, 2026 02:48

demandal25 changed the title ~~Enable torch compile~~ Enable torch.compile under a flag Apr 3, 2026

Copilot started reviewing on behalf of demandal25 April 3, 2026 02:49 View session

Copilot AI reviewed Apr 3, 2026

View reviewed changes

Comment thread flashinfer/utils.py

Comment thread scripts/verify_enable_torch_compile.py Outdated

demandal25 requested a review from diptorupd April 3, 2026 03:03

diptorupd reviewed Apr 3, 2026

View reviewed changes

Comment thread flashinfer/utils.py

diptorupd reviewed Apr 3, 2026

View reviewed changes

Comment thread scripts/verify_enable_torch_compile.py Outdated

diptorupd reviewed Apr 3, 2026

View reviewed changes

Comment thread flashinfer/get_include_paths.py Outdated

demandal25 added 3 commits April 3, 2026 17:53

Revert get_include_paths.py source-tree fallback

385d9fd

Drop the repo-root include fallback per reviewer feedback: running from a source tree without an install is not a supported workflow. Made-with: Cursor

demandal25 requested a review from diptorupd April 3, 2026 18:48

Fall back to compile guard when custom_op schema inference fails

44bcefc

torch.library.custom_op cannot infer schemas for parameters like Optional[torch.Generator]. Instead of crashing, fall back to _guard_compile so these ops still get torch.compile protection. Made-with: Cursor

Copilot AI review requested due to automatic review settings April 3, 2026 19:09

Copilot started reviewing on behalf of demandal25 April 3, 2026 19:11 View session

Copilot AI reviewed Apr 3, 2026

View reviewed changes

Comment thread tests/rocm_tests/test_torch_compile_hip.py

Comment thread flashinfer/decode_rocm.py

demandal25 requested a review from Copilot April 3, 2026 19:18

Copilot started reviewing on behalf of demandal25 April 3, 2026 19:20 View session

diptorupd approved these changes Apr 3, 2026

View reviewed changes

diptorupd merged commit 2ecbe16 into AMD-Ecosystem:amd-integration Apr 3, 2026
7 of 8 checks passed

Copilot AI reviewed Apr 3, 2026

View reviewed changes

Comment thread flashinfer/utils.py

Comment thread flashinfer/prefill_rocm.py

Comment thread flashinfer/prefill_rocm.py

Comment thread flashinfer/decode_rocm.py

Comment thread flashinfer/decode_rocm.py

demandal25 deleted the enable-torch-compile branch May 12, 2026 05:10

demandal25 mentioned this pull request May 18, 2026

fix(hip): address Copilot review comments on PR #210 torch.compile flag #230

Merged

2 tasks

demandal25 mentioned this pull request May 21, 2026

docs: update README for amd-flashinfer library consumers #239

Merged

6 tasks

Navigation Menu

Uh oh!

Enable torch.compile under a flag - #210

Enable torch.compile under a flag#210
diptorupd merged 10 commits into
AMD-Ecosystem:amd-integrationfrom
demandal25:enable-torch-compile

demandal25 commented Mar 31, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

diptorupd left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

demandal25 commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

How to use it

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

diptorupd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

demandal25 commented Mar 31, 2026 •

edited

Loading