pyproject.toml

[build-system]
requires = ["setuptools>=64", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "flash-rt"
version = "0.1.0"
description = "FlashRT — realtime CUDA inference engine for small-batch, latency-sensitive AI workloads. Ships VLA frontends for Pi0, Pi0.5, Pi0-FAST, GROOT N1.6/N1.7."
readme = "README.md"
license = {text = "Apache-2.0"}
requires-python = ">=3.10"
dependencies = [
    "numpy",
    "pyyaml",
]

[project.optional-dependencies]
torch = ["torch", "safetensors", "sentencepiece"]
jax = ["jax", "jaxlib", "ml_dtypes", "orbax-checkpoint", "flax"]
server = ["fastapi", "uvicorn"]
eval = ["opencv-python", "tqdm"]
# MiniMax-Remover NVFP4 video-inpainting pipeline. Optional: only needed to
# construct/run flash_rt.models.minimax_remover (importing the package needs
# none of these). diffusers + einops are the runtime model deps; scipy is used
# by the quickstart mask helpers; sageattention is the default attention backend
# (FLASHRT_ATTN_MODE=sage_fp8) -- switch to FLASHRT_ATTN_MODE=fa2 to drop it.
# triton ships with torch on CUDA.
minimax-remover = ["diffusers", "einops", "scipy", "sageattention"]
# Optional. Required only for the legacy upstream attention path
# (FVK_RTX_FA2=0 / FVK_RTX_FA2_SITES=...) and the GROOT N1.6/N1.7
# backend. Default RTX Pi0 / Pi0.5 inference uses the vendored
# flash_rt_fa2.so and does NOT need this. If pip cannot find a
# prebuilt wheel for your torch+CUDA+Python combo, grab one from
# https://github.com/Dao-AILab/flash-attention/releases — building
# the sdist takes 30+ minutes on cold cloud images.
flash-attn = ["flash-attn"]
# Thor-only optional fast path: FlashAttention-4 (CuTe-DSL) denoise/prefix
# attention for LingBot-VLA on Thor (sm_110). The forward / SM100-only source
# is vendored (privately, as `flashrt_fa4`) at csrc/attention/flash_attn_4_src
# (no flash-attn wheel needed); these are its runtime deps. Without them,
# LingBot silently falls back to the fmha path (correct, ~+18ms@25). The loader
# (flash_rt/hardware/thor/fa4_backend.py) sets CUTE_DSL_ARCH=sm_101a.
# Requires an editable / source-tree install (pip install -e). pip install ".[thor-fa4]".
thor-fa4 = ["nvidia-cutlass-dsl==4.5.1", "quack-kernels==0.4.1"]
all = ["flash-rt[torch,jax,server,eval,minimax-remover]"]

[tool.setuptools.packages.find]
include = ["flash_rt*"]

[tool.setuptools.package-data]
flash_rt = ["*.so", "*.yaml", "*.model"]