Instructions to use replicate/flashinfer-draft with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Kernels
How to use replicate/flashinfer-draft with Kernels:
# !pip install kernels from kernels import get_kernel kernel = get_kernel("replicate/flashinfer-draft") - Notebooks
- Google Colab
- Kaggle
| [general] | |
| name = "flashinfer" | |
| universal = false | |
| [torch] | |
| src = [ | |
| "torch-ext/torch_binding.cpp", | |
| "torch-ext/torch_binding.h" | |
| ] | |
| [kernel.flashinfer] | |
| backend = "cuda" | |
| cuda-capabilities = [ | |
| # "7.0", | |
| # "7.2", | |
| "7.5", | |
| "8.0", | |
| "8.6", | |
| "8.7", | |
| "8.9", | |
| "9.0", | |
| # "10.0", | |
| # "10.1", | |
| # "11.8", | |
| # "12.0" | |
| ] | |
| cuda-flags = [ | |
| "-O3", | |
| "-std=c++17", | |
| "--use_fast_math", | |
| "--expt-relaxed-constexpr", | |
| "--expt-extended-lambda", | |
| "-DFLASHINFER_ENABLE_F16", | |
| "-DFLASHINFER_ENABLE_BF16", | |
| "-DFLASHINFER_ENABLE_FP8_E4M3", | |
| "-DFLASHINFER_ENABLE_FP8_E5M2", | |
| "-DNDEBUG", | |
| ] | |
| cxx-flags = [ | |
| "-DFLASHINFER_ENABLE_F16", | |
| "-DFLASHINFER_ENABLE_BF16", | |
| "-DFLASHINFER_ENABLE_FP8_E4M3", | |
| "-DFLASHINFER_ENABLE_FP8_E5M2", | |
| ] | |
| include = [ "include" ] | |
| depends = ["torch", "cutlass_3_8"] | |
| src = [ | |
| "include/flashinfer/page.cuh", | |
| "include/flashinfer/arch_condition.h", | |
| "include/flashinfer/cubin_loader.h", | |
| "include/flashinfer/permuted_smem.cuh", | |
| "include/flashinfer/exception.h", | |
| "include/flashinfer/trtllm/common.h", | |
| "include/flashinfer/trtllm/fmha/fmhaRunnerParams.h", | |
| "include/flashinfer/trtllm/fmha/decoder_impl_common.h", | |
| "include/flashinfer/trtllm/fmha/decoder_params.h", | |
| "include/flashinfer/trtllm/fmha/kernelParams.h", | |
| "include/flashinfer/trtllm/fmha/gen_kernel_launcher.cuh", | |
| "include/flashinfer/trtllm/fmha/fmhaKernels.cuh", | |
| "include/flashinfer/trtllm/fmha/cubin/kernelMetaInfo.h", | |
| "include/flashinfer/trtllm/fmha/fmhaRunner.cuh", | |
| "include/flashinfer/trtllm/common/cudaTypeUtils.cuh", | |
| "include/flashinfer/trtllm/common/cudaUtils.h", | |
| "include/flashinfer/trtllm/common/cudaBf16Wrapper.h", | |
| "include/flashinfer/trtllm/common/cudaFp8Utils.h", | |
| "include/flashinfer/trtllm/common/cudaBf16Fallbacks.cuh", | |
| "include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh", | |
| "include/flashinfer/trtllm/fused_moe/DevKernel.h", | |
| "include/flashinfer/trtllm/fused_moe/RoutingKernel.h", | |
| "include/flashinfer/trtllm/fused_moe/runner.h", | |
| "include/flashinfer/trtllm/fused_moe/IntFastDiv.h", | |
| "include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh", | |
| "include/flashinfer/trtllm/batched_gemm/KernelRunner.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelMetaInfo.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h", | |
| "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/MmaDecl.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/SfLayoutDecl.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/DtypeDecl.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CommonUtils.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CudaKernelLauncher.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/TmaDescriptor.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelTraits.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmInterface.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelParams.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/Enums.h", | |
| "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmOptions.h", | |
| "include/flashinfer/page.cuh", | |
| "include/flashinfer/vec_dtypes.cuh", | |
| "include/flashinfer/sampling.cuh", | |
| "include/flashinfer/logging.h", | |
| "include/flashinfer/fp16.h", | |
| "include/flashinfer/attention_impl.cuh", | |
| "include/flashinfer/allocator.h", | |
| "include/flashinfer/utils.cuh", | |
| "include/flashinfer/fastdiv.cuh", | |
| "include/flashinfer/norm.cuh", | |
| "include/flashinfer/layout.cuh", | |
| "include/flashinfer/math.cuh", | |
| "include/flashinfer/gemm/group_gemm_sm90.cuh", | |
| "include/flashinfer/gemm/group_gemv.cuh", | |
| "include/flashinfer/gemm/fp8_gemm_cutlass_template.h", | |
| "include/flashinfer/gemm/group_gemm_mxfp4_groupwise_sm100.cuh", | |
| "include/flashinfer/gemm/fp4_gemm_cutlass.h", | |
| "include/flashinfer/gemm/fp4_gemm_cutlass_template.h", | |
| "include/flashinfer/gemm/cutlass_gemm_configs.h", | |
| "include/flashinfer/gemm/fp4_gemm_template_sm100.h", | |
| "include/flashinfer/gemm/group_gemm_fp8_groupwise_sm100.cuh", | |
| "include/flashinfer/gemm/gemm_groupwise_sm100.cuh", | |
| "include/flashinfer/gemm/group_gemm_lora.cuh", | |
| "include/flashinfer/gemm/group_gemm.cuh", | |
| "include/flashinfer/gemm/fp8_gemm_cutlass.h", | |
| "include/flashinfer/gemm/fp8_gemm_template_sm100.h", | |
| "include/flashinfer/gemm/bmm_fp8.cuh", | |
| "include/flashinfer/activation.cuh", | |
| "include/flashinfer/semaphore_utils.cuh", | |
| "include/flashinfer/quantization.cuh", | |
| "include/flashinfer/mma.cuh", | |
| "include/flashinfer/attention/variant_helper.cuh", | |
| "include/flashinfer/attention/hopper.cuh", | |
| "include/flashinfer/attention/prefill.cuh", | |
| "include/flashinfer/attention/persistent_template.cuh", | |
| "include/flashinfer/attention/mla_hopper.cuh", | |
| "include/flashinfer/attention/pod.cuh", | |
| "include/flashinfer/attention/variants.cuh", | |
| "include/flashinfer/attention/decode.cuh", | |
| "include/flashinfer/attention/heap.h", | |
| "include/flashinfer/attention/state.cuh", | |
| "include/flashinfer/attention/mask.cuh", | |
| "include/flashinfer/attention/default_prefill_params.cuh", | |
| "include/flashinfer/attention/cutlass_mla.cuh", | |
| "include/flashinfer/attention/scheduler.cuh", | |
| "include/flashinfer/attention/blackwell/kernel/fmha_options.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/gather_tensor.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/fmha_tile_scheduler.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_tma_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/sm100_mla_tile_scheduler.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_reduction.hpp", | |
| "include/flashinfer/attention/blackwell/plan.cuh", | |
| "include/flashinfer/attention/blackwell/common/pow_2.hpp", | |
| "include/flashinfer/attention/blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/collective/sm100_fmha_load_cpasync_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_epilogue_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp", | |
| "include/flashinfer/attention/blackwell/collective/fmha_fusion.hpp", | |
| "include/flashinfer/attention/blackwell/collective/fmha_common.hpp", | |
| "include/flashinfer/attention/blackwell/device/sm100_mla.hpp", | |
| "include/flashinfer/attention/blackwell/device/fmha.hpp", | |
| "include/flashinfer/attention/blackwell/fmha_cutlass_sm100.cuh", | |
| "include/flashinfer/attention/mla.cuh", | |
| "include/flashinfer/attention/decode_mla_cute_sm80.cuh", | |
| "include/flashinfer/attention/cascade.cuh", | |
| "include/flashinfer/attention/hopper/variant_helper.cuh", | |
| "include/flashinfer/attention/hopper/default_params.cuh", | |
| "include/flashinfer/attention/hopper/attention_updater.cuh", | |
| "include/flashinfer/attention/hopper/epilogue.cuh", | |
| "include/flashinfer/attention/hopper/variants.cuh", | |
| "include/flashinfer/attention/hopper/mainloop.cuh", | |
| "include/flashinfer/attention/hopper/block_sparse_gather.cuh", | |
| "include/flashinfer/attention/hopper/utils.cuh", | |
| "include/flashinfer/attention/hopper/prefill_sm90.cuh", | |
| "include/flashinfer/attention/hopper/named_barrier.cuh", | |
| "include/flashinfer/attention/hopper/tile_scheduler.cuh", | |
| "include/flashinfer/attention/hopper/sparse_mainloop.cuh", | |
| "include/flashinfer/attention/hopper/kernel_traits.cuh", | |
| "include/flashinfer/attention/hopper/mainloop_mma.cuh", | |
| "include/flashinfer/attention/hopper/quantization/epilogue.cuh", | |
| "include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh", | |
| "include/flashinfer/attention/hopper/quantization/mainloop_load.cuh", | |
| "include/flashinfer/attention/hopper/quantization/mainloop_sparse_load.cuh", | |
| "include/flashinfer/attention/hopper/quantization/kernel_traits.cuh", | |
| "include/flashinfer/attention/hopper/quantization/mainloop_mma.cuh", | |
| "include/flashinfer/attention/default_decode_params.cuh", | |
| "include/flashinfer/attention/mla_params.cuh", | |
| "include/flashinfer/attention/persistent.cuh", | |
| "include/flashinfer/pos_enc.cuh", | |
| "include/flashinfer/cutlass_utils.cuh", | |
| "include/flashinfer/comm/trtllm_moe_allreduce_fusion.cuh", | |
| "include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh", | |
| "include/flashinfer/comm/vllm_custom_all_reduce.cuh", | |
| "include/flashinfer/comm/trtllm_allreduce.cuh", | |
| "include/flashinfer/comm/trtllm_allreduce_fusion.cuh", | |
| "include/flashinfer/comm/trtllm_alltoall.cuh", | |
| "include/flashinfer/frag_layout_swizzle.cuh", | |
| "include/flashinfer/profiler.cuh", | |
| "include/flashinfer/cp_async.cuh", | |
| "include/pytorch_conversion_utils.h", | |
| "include/pytorch_extension_utils.h", | |
| # # batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc", | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu", | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu", | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu", | |
| # # batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc", | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu", | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu", | |
| # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu", | |
| "csrc/generated/gelu_and_mul.cu", | |
| "csrc/generated/gelu_tanh_and_mul.cu", | |
| "csrc/generated/silu_and_mul.cu", | |
| ] |