Skip to content

Enable torch.compile under a flag #210

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
diptorupd merged 10 commits into from
Apr 3, 2026
Merged

Enable torch.compile under a flag #210

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
4e3bb90
Enable torch.compile under a flag
demandal25 Mar 31, 2026
b1bc427
Fix to enable torch.compile
demandal25 Apr 1, 2026
f4da9a9
Revert cuda related changes
demandal25 Apr 2, 2026
4becdb0
Refactor the changelist
demandal25 Apr 2, 2026
c9ef604
Fix sampling
demandal25 Apr 2, 2026
c7b0ef2
error handling when torch.compile called witout env
demandal25 Apr 2, 2026
0afec0b
Address PR review: fix custom-op registration and convert script to p…
demandal25 Apr 3, 2026
385d9fd
Revert get_include_paths.py source-tree fallback
demandal25 Apr 3, 2026
9737731
Fix torch.compile test: use correct device and relax error matching
demandal25 Apr 3, 2026
44bcefc
Fall back to compile guard when custom_op schema inference fails
demandal25 Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions flashinfer/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@
prepare_low_latency_gemm_weights as prepare_low_latency_gemm_weights,
)
from .utils import next_positive_power_of_2 as next_positive_power_of_2
from .utils import use_torch_custom_ops_enabled as use_torch_custom_ops_enabled
from .xqa import xqa as xqa
from .xqa import xqa_mla as xqa_mla
elif IS_HIP:
Expand Down Expand Up @@ -259,6 +260,7 @@
sys.modules["flashinfer.decode"] = sys.modules["flashinfer.decode_rocm"]

from .utils import next_positive_power_of_2 as next_positive_power_of_2
from .utils import use_torch_custom_ops_enabled as use_torch_custom_ops_enabled
else:
# CPU-only torch (no CUDA or HIP)
raise RuntimeError(
Expand Down
22 changes: 16 additions & 6 deletions flashinfer/decode_rocm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
canonicalize_torch_dtype,
device_support_pdl,
is_float8,
plan_info_vec_as_tensor,
register_custom_op,
register_fake_op,
)
Expand Down Expand Up @@ -137,7 +138,7 @@ def get_batch_decode_jit_module(module_name: str, jit_module: Any):
def run_batch_decode(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: Optional[torch.Tensor],
paged_v_cache: Optional[torch.Tensor],
Expand Down Expand Up @@ -173,7 +174,7 @@ def run_batch_decode(
def _fake_run_batch_decode(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: Optional[torch.Tensor],
paged_v_cache: Optional[torch.Tensor],
Expand Down Expand Up @@ -218,7 +219,7 @@ def get_batch_decode_module(*args):
def run_batch_decode(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: Optional[torch.Tensor],
paged_v_cache: Optional[torch.Tensor],
Expand Down Expand Up @@ -262,7 +263,7 @@ def run_batch_decode(
def _fake_run_batch_decode(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: Optional[torch.Tensor],
paged_v_cache: Optional[torch.Tensor],
Expand Down Expand Up @@ -972,6 +973,9 @@ def plan(
head_dim,
False, # causal
)
self._plan_info = plan_info_vec_as_tensor(
self._plan_info, device=self._float_workspace_buffer.device
)
Comment thread
demandal25 marked this conversation as resolved.
else:
if self._jit_module is not None:
self._cached_module = self._jit_module
Expand Down Expand Up @@ -1005,6 +1009,9 @@ def plan(
torch.empty(0, dtype=q_data_type),
torch.empty(0, dtype=kv_data_type),
)
self._plan_info = plan_info_vec_as_tensor(
self._plan_info, device=self._float_workspace_buffer.device
)
Comment thread
demandal25 marked this conversation as resolved.

self._pos_encoding_mode = pos_encoding_mode
self._window_left = window_left
Expand Down Expand Up @@ -1230,13 +1237,16 @@ def run(
rope_theta,
0, # token_pos_in_items_len
self._workspace_size,
paged_kv_cache,
self._num_qo_heads,
self._num_kv_heads,
self._block_tables,
self._kv_lens_buffer,
page_size,
None, # max_q_len (decode: single token)
self._max_kv_len,
None, # batch_size
None, # cum_seq_lens_q
None, # cum_seq_lens_kv
sinks,
]

Expand All @@ -1248,7 +1258,7 @@ def run(
run_args = [
self._float_workspace_buffer,
self._int_workspace_buffer,
self._plan_info,
plan_info,
q,
Comment thread
demandal25 marked this conversation as resolved.
k_cache,
v_cache,
Expand Down
24 changes: 18 additions & 6 deletions flashinfer/get_include_paths.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,14 @@
#
# SPDX - License - Identifier : Apache - 2.0

import os
import pathlib

_INCLUDE_MARKER = pathlib.Path("flashinfer") / "attention" / "generic" / "prefill.cuh"


def _include_dir_has_headers(include_root: pathlib.Path) -> bool:
return (include_root / _INCLUDE_MARKER).is_file()


def _get_package_root_dir():
"""Return the root directory of the flashinfer package.
Expand All @@ -26,17 +31,24 @@ def _get_package_root_dir():
def get_include():
"""Return the directory containing the header files needed by the JIT.

The `include` dir in the splatlib/flashinfer directory contains the header
files needed by the JIT to compile the C++ code. The include path contains
all flashinfer, cutlass, and Cute headers and any future dependencies.
Prefer ``<package>/flashinfer/include`` (wheel install or editable symlink from
CMake). If that directory does not contain the JIT headers yet, fall back to
``<repo>/include`` so running from a source tree with ``PYTHONPATH`` works
without running the editable install symlink step.

Returns
-------
include_dir : str
Path to include and Cutlass header files.
"""
include_dir = os.path.join(_get_package_root_dir(), "include")
return str(include_dir)
package_dir = pathlib.Path(_get_package_root_dir()).resolve()
Comment thread
demandal25 marked this conversation as resolved.
Outdated
pkg_include = package_dir / "include"
repo_include = package_dir.parent / "include"
if _include_dir_has_headers(pkg_include):
return str(pkg_include)
if _include_dir_has_headers(repo_include):
return str(repo_include)
return str(pkg_include)


def get_csrc_dir():
Expand Down
36 changes: 21 additions & 15 deletions flashinfer/prefill_rocm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
determine_attention_backend,
device_support_pdl,
is_float8,
plan_info_vec_as_tensor,
register_custom_op,
register_fake_op,
)
Expand Down Expand Up @@ -288,7 +289,6 @@ def run_single_prefill(
1.0 / rope_scale, # rope_rcp_scale
1.0 / rope_theta, # rope_rcp_theta
)
return o

@register_fake_op(f"flashinfer::{uri}_run")
def _fake_run_single_prefill(
Expand All @@ -305,6 +305,9 @@ def _fake_run_single_prefill(
maybe_alibi_slopes: Optional[torch.Tensor],
logits_soft_cap: float,
sm_scale: float,
scale_q: Optional[torch.Tensor],
scale_k: Optional[torch.Tensor],
scale_v: Optional[torch.Tensor],
rope_scale: float,
rope_theta: float,
) -> None:
Expand Down Expand Up @@ -525,7 +528,7 @@ def get_batch_prefill_module(backend, *args):
def ragged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
Expand Down Expand Up @@ -581,13 +584,11 @@ def ragged_run(
# token_pos_in_items_len, # Not supported by HIP FA2 kernels
)

return o

@register_fake_op(f"flashinfer::{uri}_ragged_run")
def _fake_ragged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
Expand Down Expand Up @@ -629,7 +630,7 @@ def _fake_ragged_run(
def paged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: torch.Tensor,
paged_v_cache: torch.Tensor,
Expand Down Expand Up @@ -705,13 +706,11 @@ def paged_run(
# token_pos_in_items_len, # Not supported by HIP FA2 kernels
)

return o

@register_fake_op(f"flashinfer::{uri}_paged_run")
def _fake_paged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: torch.Tensor,
paged_v_cache: torch.Tensor,
Expand Down Expand Up @@ -780,7 +779,7 @@ def get_batch_prefill_jit_module(module_name: str, jit_module: Any):
def ragged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
Expand Down Expand Up @@ -814,7 +813,7 @@ def ragged_run(
def _fake_ragged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
Expand Down Expand Up @@ -844,7 +843,7 @@ def _fake_ragged_run(
def paged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: torch.Tensor,
paged_v_cache: torch.Tensor,
Expand Down Expand Up @@ -882,7 +881,7 @@ def paged_run(
def _fake_paged_run(
float_workspace_buffer: torch.Tensor,
int_workspace_buffer: torch.Tensor,
plan_info_vec: List[int],
plan_info_vec: torch.Tensor,
q: torch.Tensor,
paged_k_cache: torch.Tensor,
paged_v_cache: torch.Tensor,
Expand Down Expand Up @@ -1525,7 +1524,7 @@ def __init__(
self._mask_indptr_buf = mask_indptr_buf
self._max_total_num_rows = None
self._backend = backend
self._plan_info = None
self._plan_info: Optional[torch.Tensor] = None
self._cached_module = None
self._seq_lens_kv = None
self._seq_lens_q = None
Expand Down Expand Up @@ -1894,6 +1893,9 @@ def plan(
head_dim_vo,
causal,
)
self._plan_info = plan_info_vec_as_tensor(
self._plan_info, device=self._float_workspace_buffer.device
Comment thread
demandal25 marked this conversation as resolved.
)

self._causal = causal
self._pos_encoding_mode = pos_encoding_mode
Expand Down Expand Up @@ -2504,6 +2506,7 @@ def __init__(
self._mask_indptr_buf = mask_indptr_buf
self._max_total_num_rows = None
self._backend = backend
self._plan_info: Optional[torch.Tensor] = None
self._cached_module = None

@property
Expand Down Expand Up @@ -2784,6 +2787,9 @@ def plan(
head_dim_vo,
causal,
)
self._plan_info = plan_info_vec_as_tensor(
self._plan_info, device=self._float_workspace_buffer.device
)
Comment thread
demandal25 marked this conversation as resolved.

self._causal: bool = causal
self._pos_encoding_mode = pos_encoding_mode
Expand Down Expand Up @@ -2920,7 +2926,7 @@ def run(
run_args = [
self._float_workspace_buffer,
self._int_workspace_buffer,
self._plan_info, # type: ignore[has-type]
self._plan_info,
q,
k,
v,
Expand Down
Loading
Loading