perf: reduce torch.library dispatch overhead (#968)

This pr fixes #960,
#764 refactors the codebase to use torch.library which introduces some
cpu-side overhead when operators are not captured by CUDAGraph.

As suggested by @youkaichao , we can bypass the pytorch dispatcher by
`torch.ops.namespace.op_name.default`.

Co-authored-by: Kaichao You <[email protected]>

Author: yzh119

flashinfer/activation.py

@@ -66,7 +66,7 @@ def get_act_and_mul_module(act_func_name: str):
 
         # torch library for act_and_mul
         fname = f"{act_func_name}_and_mul"
-        fn = getattr(module, fname)
+        fn = getattr(module, fname).default
 
         @register_custom_op(f"flashinfer::{fname}", mutates_args=("out",))
         def _act_and_mul(

flashinfer/cascade.py

@@ -98,7 +98,7 @@ def merge_state(
         s_b = s_b.to(torch.float32)
         v_merged = torch.empty_like(v_a)
         s_merged = torch.empty_like(s_a)
-        get_cascade_module().merge_state(
+        get_cascade_module().merge_state.default(
             v_a, s_a, v_b, s_b, v_merged, s_merged, get_cuda_stream(device)
         )
         return v_merged, s_merged
@@ -160,7 +160,7 @@ def merge_state_in_place(
     with v.device as device:  # device guard
         s = s.to(torch.float32)
         s_other = s_other.to(torch.float32)
-        get_cascade_module().merge_state_in_place(
+        get_cascade_module().merge_state_in_place.default(
             v, s, v_other, s_other, mask, get_cuda_stream(device)
         )
 
@@ -221,7 +221,7 @@ def merge_states(v: torch.Tensor, s: torch.Tensor) -> Tuple[torch.Tensor, torch.
             seq_len, num_heads, head_dim, dtype=v.dtype, device=device
         )
         s_merged = torch.empty(seq_len, num_heads, dtype=torch.float32, device=device)
-        get_cascade_module().merge_states(
+        get_cascade_module().merge_states.default(
             v, s, v_merged, s_merged, get_cuda_stream(device)
         )
         return v_merged, s_merged

flashinfer/decode.py

@@ -69,9 +69,9 @@ def get_single_decode_module(*args):
         if has_prebuilt_ops and uri in prebuilt_ops_uri:
             _kernels = torch.ops.flashinfer_kernels
 
-            run_func = _kernels.single_decode_with_kv_cache
+            run_func = _kernels.single_decode_with_kv_cache.default
         else:
-            run_func = gen_single_decode_module(*args).run
+            run_func = gen_single_decode_module(*args).run.default
 
         # torch library for single_decode_with_kv_cache
 
@@ -134,8 +134,8 @@ def get_batch_decode_jit_module(module_name: str, jit_module: Any):
     if module_name in _batch_decode_jit_modules:
         return _batch_decode_jit_modules[module_name]
 
-    plan_func = jit_module.plan
-    run_func = jit_module.run
+    plan_func = jit_module.plan.default
+    run_func = jit_module.run.default
 
     @register_custom_op(
         f"flashinfer::{module_name}_run",
@@ -216,12 +216,12 @@ def get_batch_decode_module(*args):
         if has_prebuilt_ops and uri in prebuilt_ops_uri:
             _kernels = torch.ops.flashinfer_kernels
 
-            plan_func = _kernels.batch_decode_with_paged_kv_cache_plan
-            run_func = _kernels.batch_decode_with_paged_kv_cache_run
+            plan_func = _kernels.batch_decode_with_paged_kv_cache_plan.default
+            run_func = _kernels.batch_decode_with_paged_kv_cache_run.default
         else:
             mod = gen_batch_decode_module(*args)
-            plan_func = mod.plan
-            run_func = mod.run
+            plan_func = mod.plan.default
+            run_func = mod.run.default
 
         # torch library for batch_decode_with_paged_kv_cache_run
 
@@ -327,7 +327,7 @@ def single_decode_with_kv_cache_with_jit_module(
             "single_decode_with_kv_cache_tmp", 32 * 1024 * 1024, device
         )
         o = torch.empty_like(q)
-        jit_module.run(
+        jit_module.run.default(
             q,
             k,
             v,

flashinfer/gemm.py

@@ -66,7 +66,7 @@ def bmm_fp8(
         ) -> None:
             with A.device as device:
                 cublas_handle = torch.cuda.current_blas_handle()
-                module.bmm_fp8(
+                module.bmm_fp8.default(
                     A,
                     B,
                     D,
@@ -105,7 +105,7 @@ def cutlass_segment_gemm(
             weight_column_major: bool,
         ) -> None:
             with x_data.device as device:
-                module.cutlass_segment_gemm(
+                module.cutlass_segment_gemm.default(
                     workspace_buffer,
                     all_problems,
                     x_data,
@@ -182,7 +182,7 @@ def cutlass_segment_gemm_sm90(
             weight_column_major: bool,
         ) -> None:
             with x_data.device as device:
-                module.cutlass_segment_gemm_sm90(
+                module.cutlass_segment_gemm_sm90.default(
                     workspace_buffer,
                     int_workspace_buffer,
                     all_problems,

flashinfer/mla.py

@@ -254,7 +254,7 @@ def plan(
         self._use_profiler = use_profiler
 
         with self.device as device:
-            self._plan_info = self._cached_module.plan(
+            self._plan_info = self._cached_module.plan.default(
                 self._float_workspace_buffer,
                 self._int_workspace_buffer,
                 self._pin_memory_int_workspace_buffer,
@@ -349,7 +349,7 @@ def run(
                         lse, q_nope.shape[:2], torch.float32, q_nope.device, "lse"
                     )
             profiler_args = (profiler_buffer,) if self._use_profiler else ()
-            self._cached_module.run(
+            self._cached_module.run.default(
                 self._float_workspace_buffer,
                 self._int_workspace_buffer,
                 self._plan_info,

flashinfer/norm.py

@@ -87,7 +87,7 @@ def _rmsnorm(
     enable_pdl: bool,
 ) -> None:
     with input.device as device:  # device guard
-        get_norm_module().rmsnorm(
+        get_norm_module().rmsnorm.default(
             out, input, weight, eps, enable_pdl, get_cuda_stream(device)
         )
 
@@ -134,7 +134,7 @@ def fused_add_rmsnorm(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
     """
     with input.device as device:  # device guard
-        get_norm_module().fused_add_rmsnorm(
+        get_norm_module().fused_add_rmsnorm.default(
             input, residual, weight, eps, enable_pdl, get_cuda_stream(device)
         )
 
@@ -195,7 +195,7 @@ def _gemma_rmsnorm(
     enable_pdl: bool,
 ) -> None:
     with input.device as device:  # device guard
-        get_norm_module().gemma_rmsnorm(
+        get_norm_module().gemma_rmsnorm.default(
             out, input, weight, eps, enable_pdl, get_cuda_stream(device)
         )
 
@@ -244,7 +244,7 @@ def gemma_fused_add_rmsnorm(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
     """
     with input.device as device:
-        get_norm_module().gemma_fused_add_rmsnorm(
+        get_norm_module().gemma_fused_add_rmsnorm.default(
             input, residual, weight, eps, enable_pdl, get_cuda_stream(device)
         )
 

flashinfer/page.py

@@ -72,7 +72,7 @@ def block_sparse_indices_to_vector_sparse_offsets(
         assert vector_sparse_indptr.dtype == torch.int32
         assert kv_lens.dtype == torch.int32
         batch_size = block_sparse_indptr.size(0) - 1
-        get_page_module().block_sparse_indices_to_vector_sparse_offsets(
+        get_page_module().block_sparse_indices_to_vector_sparse_offsets.default(
             block_sparse_indices,
             block_sparse_indptr,
             vector_sparse_offsets,
@@ -108,7 +108,7 @@ def _append_paged_mla_kv_cache_kernel(
         kv_indices = kv_indices.int()
         kv_indptr = kv_indptr.int()
         kv_last_page_len = kv_last_page_len.int()
-        get_page_module().append_paged_mla_kv_cache(
+        get_page_module().append_paged_mla_kv_cache.default(
             append_ckv,
             append_kpe,
             batch_indices,
@@ -144,7 +144,7 @@ def _append_paged_kv_cache_kernel(
         kv_indices = kv_indices.int()
         kv_indptr = kv_indptr.int()
         kv_last_page_len = kv_last_page_len.int()
-        get_page_module().append_paged_kv_cache(
+        get_page_module().append_paged_kv_cache.default(
             append_key,
             append_value,
             batch_indices,

flashinfer/pod.py

@@ -67,9 +67,9 @@ def get_pod_module(*args):
             _kernels = torch.ops.flashinfer_kernels
             # torch library for pod_with_kv_cache
             # No tensor deprecated due to poor performance. Just use tensor cores for both.
-            run_tensor = _kernels.pod_with_kv_cache_tensor
+            run_tensor = _kernels.pod_with_kv_cache_tensor.default
         else:
-            run_tensor = gen_pod_module(*args).pod_with_kv_cache_tensor
+            run_tensor = gen_pod_module(*args).pod_with_kv_cache_tensor.default
         # Register the module
         _pod_modules[args] = SimpleNamespace(run_tensor=run_tensor)
     return _pod_modules[args]

flashinfer/prefill.py

@@ -73,13 +73,13 @@ def backend_module(*args):
                 if backend == "fa2":
                     _kernels = torch.ops.flashinfer_kernels
 
-                    run_func = _kernels.single_prefill_with_kv_cache
+                    run_func = _kernels.single_prefill_with_kv_cache.default
                 else:
                     _kernels_sm90 = torch.ops.flashinfer_kernels_sm90
 
-                    run_func = _kernels_sm90.single_prefill_with_kv_cache_sm90
+                    run_func = _kernels_sm90.single_prefill_with_kv_cache_sm90.default
             else:
-                run_func = gen_single_prefill_module(backend, *args).run
+                run_func = gen_single_prefill_module(backend, *args).run.default
 
             # torch library for single_prefill_with_kv_cache
 
@@ -180,24 +180,30 @@ def backend_module(*args):
                 if backend == "fa2":
                     _kernels = torch.ops.flashinfer_kernels
 
-                    plan_func = _kernels.batch_prefill_with_kv_cache_plan
-                    ragged_run_func = _kernels.batch_prefill_with_ragged_kv_cache_run
-                    paged_run_func = _kernels.batch_prefill_with_paged_kv_cache_run
+                    plan_func = _kernels.batch_prefill_with_kv_cache_plan.default
+                    ragged_run_func = (
+                        _kernels.batch_prefill_with_ragged_kv_cache_run.default
+                    )
+                    paged_run_func = (
+                        _kernels.batch_prefill_with_paged_kv_cache_run.default
+                    )
                 else:
                     _kernels_sm90 = torch.ops.flashinfer_kernels_sm90
 
-                    plan_func = _kernels_sm90.batch_prefill_with_kv_cache_sm90_plan
+                    plan_func = (
+                        _kernels_sm90.batch_prefill_with_kv_cache_sm90_plan.default
+                    )
                     ragged_run_func = (
-                        _kernels_sm90.batch_prefill_with_ragged_kv_cache_sm90_run
+                        _kernels_sm90.batch_prefill_with_ragged_kv_cache_sm90_run.default
                     )
                     paged_run_func = (
-                        _kernels_sm90.batch_prefill_with_paged_kv_cache_sm90_run
+                        _kernels_sm90.batch_prefill_with_paged_kv_cache_sm90_run.default
                     )
             else:
                 module = gen_batch_prefill_module(backend, *args)
-                plan_func = module.plan
-                ragged_run_func = module.ragged_run
-                paged_run_func = module.paged_run
+                plan_func = module.plan.default
+                ragged_run_func = module.ragged_run.default
+                paged_run_func = module.paged_run.default
 
             # torch library for ragged_run
 
@@ -437,9 +443,9 @@ def get_batch_prefill_jit_module(module_name: str, jit_module: Any):
     if module_name in _batch_prefill_jit_modules:
         return _batch_prefill_jit_modules[module_name]
 
-    plan_func = jit_module.plan
-    ragged_run_func = jit_module.ragged_run
-    paged_run_func = jit_module.paged_run
+    plan_func = jit_module.plan.default
+    ragged_run_func = jit_module.ragged_run.default
+    paged_run_func = jit_module.paged_run.default
 
     # torch library for ragged_run
     @register_custom_op(
@@ -611,7 +617,7 @@ def single_prefill_with_kv_cache_with_jit_module(
             lse = torch.empty(
                 (q.size(0), q.size(1)), dtype=torch.float32, device=device
             )
-        jit_module.run(
+        jit_module.run.default(
             q,
             k,
             v,

flashinfer/quantization.py

@@ -47,7 +47,9 @@ def _packbits(x: torch.Tensor, bitorder: str) -> torch.Tensor:
     with x.device as device:  # device guard
         x = x.to(torch.bool)
         y = torch.empty((x.size(0) + 7) // 8, dtype=torch.uint8, device=device)
-        get_quantization_module().packbits(x, bitorder, y, get_cuda_stream(device))
+        get_quantization_module().packbits.default(
+            x, bitorder, y, get_cuda_stream(device)
+        )
         return y
 
 
@@ -146,7 +148,7 @@ def segment_packbits(
         indptr = indptr.to(torch.int32)
         indptr_new = indptr_new.to(torch.int32)
         y = torch.empty(output_nnzs, dtype=torch.uint8, device=device)
-        get_quantization_module().segment_packbits(
+        get_quantization_module().segment_packbits.default(
             x, indptr, indptr_new, bitorder, y, get_cuda_stream(device)
         )
         return (

flashinfer/rope.py

@@ -58,7 +58,7 @@ def _apply_rope(
     with q.device as device:
         indptr = indptr.int()
         offsets = offsets.int()
-        get_rope_module().apply_rope(
+        get_rope_module().apply_rope.default(
             q,
             k,
             q_rope,
@@ -108,7 +108,7 @@ def _apply_llama31_rope(
     with q.device as device:
         indptr = indptr.int()
         offsets = offsets.int()
-        get_rope_module().apply_llama31_rope(
+        get_rope_module().apply_llama31_rope.default(
             q,
             k,
             q_rope,
@@ -159,7 +159,7 @@ def _apply_rope_pos_ids(
 ) -> None:
     with q.device as device:
         pos_ids = pos_ids.int()
-        get_rope_module().apply_rope_pos_ids(
+        get_rope_module().apply_rope_pos_ids.default(
             q,
             k,
             q_rope,
@@ -202,7 +202,7 @@ def _apply_rope_pos_ids_cos_sin_cache(
 ) -> None:
     with q.device as device:
         pos_ids = pos_ids.int()
-        get_rope_module().apply_rope_pos_ids_cos_sin_cache(
+        get_rope_module().apply_rope_pos_ids_cos_sin_cache.default(
             q,
             k,
             q_rope,
@@ -247,7 +247,7 @@ def _apply_llama31_rope_pos_ids(
 ) -> None:
     with q.device as device:
         pos_ids = pos_ids.int()
-        get_rope_module().apply_llama31_rope_pos_ids(
+        get_rope_module().apply_llama31_rope_pos_ids.default(
             q,
             k,
             q_rope,