Skip to content

Configurable Environment Variables

This document summarizes all configurable environment variables in ffpa-attn, which control the kernel build (extension selection, target SM list, headdim coverage, quantized-kernel windows, ccache shim) and the runtime knobs (backend gating / SDPA fallback, MMA accumulator dtype, SMEM swizzle, prefetch / persist policies, launch grid layout, persistent tuned configs). Most knobs are boolean integers parsed from 0 / 1.

Build-time environment variables

These are read once during pip install . / python setup.py build_ext and decide which translation units are generated and how nvcc is invoked.

  • FFPA_BUILD_ARCH, default "" (current device SM), Comma/semicolon/space separated list of target CUDA SM architectures. Accepts numeric SMs (e.g. "80,89,90") or aliases (maxwell, pascal, volta, turing, ampere, ada, hopper, blackwell, blackwell_geforce). When empty, falls back to the currently visible CUDA device's compute capability.
  • ENABLE_FFPA_CUDA_IMPL, default False (0), Build the optional native CUDA forward extension (ffpa_attn._C). When unset, the package builds in pure Python / Triton-only mode and large-D forward defaults to the Triton backend. Native CUDA backward is no longer generated; use Triton or SDPA backward.
  • FFPA_SKIP_CUDA_EXT, default False (0), Force-skip the CUDA extension build even if ENABLE_FFPA_CUDA_IMPL=1 is set. Used by ReadTheDocs and the check-mkdocs GitHub workflow where no nvcc is available; mkdocstrings can still import ffpa_attn to extract docstrings.
  • ENABLE_FFPA_TMA_EXT, default False (0), Compile the TMA-based kernel families (native TMA for sm90+, CuTe TMA fp16/fp8/fp4 for sm120). tools/build_fast.sh --ext all sets it; it is force-disabled when every target arch is sm<90, and the custom fp8 sm_89 family covers that case instead.
  • ENABLE_FFPA_CUTE_EXT, default False (0), Compile the CuTe C++ kernel extension (CuTe fp16 families on sm120/sm80). Needs ENABLE_FFPA_TMA_EXT=1 for the TMA-gated parts.
  • ENABLE_FFPA_FP8_SM89_EXT, default False (0), Compile the TMA-free fp8 sm_89 kernel family (persist-D D<=224, split-D 224<D<=512 with ENABLE_FFPA_QUANT_SPLIT_D=1). tools/build_fast.sh turns it on automatically when every target arch is sm<90 (L20 / 4090); requires ENABLE_FFPA_CUTE_EXT=1.
  • ENABLE_FFPA_QUANT_SPLIT_D, default False (0), Compile the split-D / M4N2 quantized attention kernels (fp8 and fp4). Off by default: --ext all builds a persist-only quantized family (fp8 D<=224, fp4 64<=D<=256), because large-D quantized attention has no speedup window over the bf16 family (PC-19 decision). Calls past the persist window raise with this env name; debug builds (ENABLE_FFPA_BUILD_DEBUG=fp8/fp4) imply it.
  • ENABLE_FFPA_F16_ACC, default False (0), Compile the fp16-input / fp16-MMA-acc kernels (CUDABackend(acc="f16")). Roughly doubles the fp16 translation-unit count; when disabled the dispatch raises a clear rebuild with ENABLE_FFPA_F16_ACC=1 error.
  • ENABLE_FFPA_CUDA_INPUT_FP16, default False (0), Compile the fp16-input CUDA kernels. Default builds are bf16-input-only (halves the forward TU count); fp16 queries then raise a "rebuild with ENABLE_FFPA_CUDA_INPUT_FP16=1" error from the Python layer.
  • ENABLE_FFPA_CUDA_MASK_FP32, default False (0), Compile the fp32-attn_mask variant translation units. Default builds keep fp16/bf16 masks only; fp32 masks are downcast to query.dtype by the Python layer (numerically equivalent, keeps every bias plan mode reachable).
  • FFPA_NVCC_THREADS, default 4, nvcc intra-TU parallelism (--threads N). With the per-headdim TU split, the outer MAX_JOBS already drives many nvcc processes in parallel, so keeping --threads small avoids oversubscription. Set to 1 to disable intra-TU threading entirely; larger values only help when MAX_JOBS is small.
  • FFPA_PTXAS_VERBOSE, default False (0), Emit ptxas verbose info (register / SMEM usage). Off by default because it produces tens of MB of log output; enable only for tuning.
  • FFPA_DEV_HEADDIMS, default "", Development-time headdim subset override. Comma/space separated list of headdims (e.g. "256,512") that replaces the full generated set for fast iteration. Empty (default) means use the full set decided by ENABLE_FFPA_ALL_HEADDIM.
  • FFPA_BUILD_STAGES, default "2,3", Explicit build stage subset with the highest priority. Comma separated list of stages (e.g. "2,3") or "all" (= stages 1..FFPA_BUILD_MAX_STAGES). Priority chain: FFPA_BUILD_STAGES > ENABLE_FFPA_ALL_STAGES (legacy, honored only when the former is unset) > default "2,3". Note that stage 1 is no longer built by default; requesting an unbuilt stage at runtime maps to the nearest built stage instead of always falling back to stage 1 (e.g. with the default "2,3", a requested stage of 1 maps to 2 and 4 maps to 3). Set FFPA_BUILD_STAGES=all (or ENABLE_FFPA_ALL_STAGES=1) to restore the previous full build. Exposed as tools/build_fast.sh --stages <csv|all>.
  • FFPA_BUILD_MAX_STAGES, default 4, Upper bound of the cp.async pipeline stages generated at build time (stages 1..N). FFPA_BUILD_STAGES=all expands to this range. Values above 4 rarely pay off.
  • ENABLE_FFPA_ALL_STAGES, legacy build-time switch, honored only when FFPA_BUILD_STAGES is unset (see above). When explicitly set, the old semantics apply: 1 generates kernels for all multi-stage variants (stages 1..FFPA_BUILD_MAX_STAGES); 0 generates only stages 1~2. When neither variable is set, the default stage set is "2,3". Reducing the stage set shortens build time at the cost of fewer schedule choices at runtime.
  • FFPA_NVCC_DEFINES, default "", Comma-separated extra -D defines injected into every nvcc invocation (debug / profiling pass-through, e.g. FFPA_NVCC_DEFINES=MY_FLAG=1).
  • FFPA_LINEINFO, default 0, When 1, add -lineinfo to nvcc so profilers (NCU / NSys) can map SASS to source. Slightly slows compilation and grows the binary.
  • ENABLE_FFPA_BUILD_DEBUG, default "", Compile per-family debug dispatch knobs (e.g. FFPA_FP8_FORCE_KERNEL, FFPA_DROPOUT_BITMAP_DISABLE, FFPA_FP4_PAD_TORCH). Comma separated subset of fp16,fp8,fp4 or all; each selected family expands to -DENABLE_FFPA_<FAM>_BUILD_DEBUG for nvcc. Empty (default) strips every debug getenv branch from the build, which also removes the FFPA_FP8_FORCE_KERNEL dual split-D template instantiation (a major fp8 compile-time cost). Note: FFPA_BIAS_TILE_DISABLE lives in a shared helper, so it becomes active for all families whenever any one family's debug macro is enabled. Exposed as tools/build_fast.sh --debug-ext <csv|all|none>.
  • ENABLE_FFPA_ALL_HEADDIM, default False (0), When 1, headdims range from 64 to 1024 with step 64 (range(64, 1024, 64)); when 0, headdims fall back to the fixed default set 64,128,192,256,320,512 (same as build_fast.sh --headdim default). Only effective on the bare setup.py/pip path; build_fast.sh always exports an explicit FFPA_DEV_HEADDIMS (which takes priority).
  • MAX_JOBS, default min(nproc, 32) via tools/build_fast.sh, Outer build parallelism passed to setuptools. The fast-build wrapper auto-caps at 32; for plain python setup.py build_ext, set it explicitly (e.g. MAX_JOBS=32).

tools/build_fast.sh-only variables

These only affect the ccache-based fast-build wrapper.

The wrapper also accepts CLI flags that map onto the variables below (flags override same-named env vars; unknown args pass through to setup.py build_ext / pip install -e):

# cuda impl + cute + tma, editable install, default headdim set, 32 jobs
bash tools/build_fast.sh --arch sm_120f --ext all --editable --headdim default --jobs 32
# equivalent env-var form:
FFPA_BUILD_ARCH=sm_120f ENABLE_FFPA_CUDA_IMPL=1 ENABLE_FFPA_CUTE_EXT=1 \
ENABLE_FFPA_TMA_EXT=1 FFPA_EDITABLE=1 MAX_JOBS=32 bash tools/build_fast.sh
flag maps to
--arch <spec> FFPA_BUILD_ARCH; full names (sm_120f, sm_89), short forms (120f) or aliases; comma-separate or repeat the flag for multiple archs
--ext all\|none\|<csv of cuda,cute,tma> ENABLE_FFPA_CUDA_IMPL / ENABLE_FFPA_CUTE_EXT / ENABLE_FFPA_TMA_EXT; all sets all three; default when omitted: ENABLE_FFPA_CUDA_IMPL=1. TMA auto-disables when every target arch is sm<90
--headdim all\|default\|<list> FFPA_DEV_HEADDIMS; all builds 64,128,192,256,320,512,768; default builds the same set minus 768; omitting the flag is equivalent to default
--stages <csv>\|all FFPA_BUILD_STAGES; stage subset, e.g. 2,3 or all (stages 1..max); takes priority over ENABLE_FFPA_ALL_STAGES; default without both: "2,3"
--debug-ext <csv>\|all\|none ENABLE_FFPA_BUILD_DEBUG; compile debug dispatch knobs for the selected families only (subset of fp16,fp8,fp4, or all); none (default) strips all debug getenv branches
--editable / --clean / --shm FFPA_EDITABLE / FFPA_CLEAN / FFPA_BUILD_IN_SHM
-j, --jobs N / --nvcc-threads N / --ptxas-verbose MAX_JOBS / FFPA_NVCC_THREADS / FFPA_PTXAS_VERBOSE
--dry-run / -h, --help print resolved env + command and exit before side effects / show usage
  • FFPA_CLEAN, default 0, When 1, removes build/, dist/, ffpa_attn.egg-info/, compiled *.so, and generated csrc/cuffpa/generated/*.{cu,h} before rebuilding.
  • FFPA_BUILD_IN_SHM, default 0, When 1, symlinks build/ into /dev/shm/ffpa-build-$USER (tmpfs) for IO-bound machines.
  • CCACHE_MAXSIZE, default 20G, Cap of the ccache storage used by the nvcc shim.
  • CCACHE_DIR, default ~/.ccache, ccache storage directory.
  • NVCC_REAL, default $CUDA_HOME/bin/nvcc, Override the path to the real nvcc (advanced; rarely needed).

Runtime kernel-selection environment variables

These are read at import / call time and gate which generated kernel template is dispatched. Most are booleans (0 / 1) unless noted.

Backend gating (small-D / fallback)

  • FFPA_CUDA_ALLOW_SMALL_D, default True (1), The CUDA backend runs small headdims (D<=256) instead of falling back to SDPA. This covers bench --cuda-impl fp8/fp4 --D 64/128 and small-D fp8/fp4 deployment; a headdim that was not compiled raises at dispatch. Set to 0 to restore the SDPA fallback for that window.
  • FFPA_TRITON_ALLOW_SMALL_D, default True (1), Same gate for the Triton backend; set to 0 to restore the SDPA fallback.
  • FFPA_CUTE_ALLOW_SMALL_D, default True (1), Same gate for the CuTeDSL backend; set to 0 to restore the SDPA fallback.
  • FFPA_FORCE_NO_SDPA_FALLBACK, default False (0), Test/dev escape hatch: when set (!= "" and != "0"), FFPAAttnMeta.fallback never reroutes to SDPA, so tests assert the selected backend's real kernel path.

Logging

  • FFPA_LOGGER_LEVEL, default INFO, Logger level for ffpa_attn (e.g. DEBUG). At DEBUG the persistent-tuned-config lookup logs each cache-hit kernel + launch config once per process (debug_once).

Persistent tuned configs (Triton)

  • FFPA_SKIP_PERSISIT_TUNED_CONFIG, default False (0), When set to 1, lookup_persistent_config() always returns None even if a device JSON exists and matches. This is useful for A/B benchmarking the persistent tuned-config path against the built-in fallback launch defaults without changing call sites.
  • FFPA_TUNED_CONFIG_DIR, default "" (package src/ffpa_attn/triton/configs), Directory to load persistent tuned configs from, e.g. a smoke-test output dir.
  • FFPA_AUTOTUNE_MAX_CONFIGS, default 0 (unlimited), Cap the number of shapes the autotune generator tunes, for fast development smoke tests.

Backend-specific behavior

  • FFPA_FP8_SM89_FORCE, default False (0), Force the fp8 sm_89 (cp.async) kernel family even on sm120+. The equivalent per-call switch is CUDABackend(force_fp8_sm89=True); both simply bypass the automatic major < 12 selection.
  • FFPA_SM89_PERSIST_KVCFG, default "64,1", fp8 sm_89 persist-D kernel geometry "kBc,stages" (kBc in {64,128}, stages in {1,2,3}). The default 64,1 keeps a 16KB smem footprint so two CTAs stay resident per SM.
  • FFPA_CUTE_SM80_FWD_USE_PERSIST_Q, default False (0), CuTeDSL sm80 forward: prefer the persist-Q variant.
  • FFPA_TRITON_BWD_FUSE_DKDVDQ, default False (0), Triton backward: fuse dK/dV/dQ into one kernel (atomic_add-based) when the shape allows; off by default (the non-fused path avoids atomics).

Debug-build-only dispatch switches

These exist only in builds with the matching ENABLE_FFPA_BUILD_DEBUG family compiled (see the build-time section); see the FP8 / FP4 guide for usage notes.

  • FFPA_FP8_FORCE_KERNEL (split_d / m4n2), force one fp8 split-D family for A/B (224 < D <= 1024).
  • FFPA_FP8_PQUANT_PER_ROW, per-row P quantization (full-scale; mutually exclusive with lazy rescale).
  • FFPA_FP4_PAD_TORCH, materialize the fp4 head-dim pad with torch instead of the fused path.
  • FFPA_BIAS_TILE_DISABLE / FFPA_BIAS_TILE_KEEP / FFPA_BIAS_ROWVEC_DISABLE, attn_bias IO-path A/B switches (tile modes / row-vector inline loads).
  • FFPA_DROPOUT_BITMAP_DISABLE, disable the fp16 dropout smem-bitmap path (fall back to inline Philox).

MMA accumulator dtype

  • ENABLE_FFPA_FORCE_QK_F16, default False (0), Force Q@K^T MMA accumulator to FP16 within the FFPA Acc-F32 kernels. Enables the mixed mode Q@K^T MMA Acc F16 + P@V MMA Acc F32.
  • ENABLE_FFPA_FORCE_PV_F16, default False (0), Force P@V MMA accumulator to FP16 within the FFPA Acc-F32 kernels. Enables the mixed mode Q@K^T MMA Acc F32 + P@V MMA Acc F16.

Prefetch & SMEM-share policies

  • ENABLE_FFPA_PREFETCH_QKV, default True (1), Prefetch QKV at the appropriate time point. Typical boost is 5%~10%.
  • ENABLE_FFPA_QKV_SMEM_SHARE, default False (0), Use a shared QKV SMEM policy. Off by default because separate buffers overlap better with MMA / g2s; turn on for low-SRAM devices.

SMEM swizzle (vs. padding) for Q / K / V

  • ENABLE_FFPA_SMEM_SWIZZLE_Q, default True (1), True: bank-conflict-free Q SMEM via swizzle. False: bank-conflict-free via padding.
  • ENABLE_FFPA_SMEM_SWIZZLE_K, default True (1), Same as above for K SMEM.
  • ENABLE_FFPA_SMEM_SWIZZLE_V, default True (1), Same as above for V SMEM.

Persistent g2s / s2r loads

  • ENABLE_FFPA_PERSIST_Q_G2S, default True (1), Persistently keep Q in SMEM via g2s for headdim <= 320. Trades more SRAM for fewer global loads while keeping register usage stable.
  • ENABLE_FFPA_PERSIST_KV_G2S, default True (1), Persistently keep KV in SMEM via g2s for headdim <= 256. When enabled, FFPA auto-uses the FlashAttention attention-level tiling for headdim <= 256 and the FFPA fine-grained MMA-level tiling for headdim > 256.
  • ENABLE_FFPA_PERSIST_Q_S2R, default False (0), Persistently load Q s2r for headdim < 512 to reduce Q g2s/s2r IO while preserving O(1) SRAM complexity. Adds register pressure as headdim grows; weigh register usage vs. IO reduction before enabling.
  • ENABLE_FFPA_PERSIST_V_S2R, default True (1), Persistently load V s2r for the small-d kernel only (more registers).

Pipelining and launch grid layout

  • ENABLE_FFPA_REGISTERS_PIPE_KV, default False (0), Use register-level ping-pong double buffers for ldmatrix / MMA computation overlap.
  • ENABLE_FFPA_LAUNCH_GRID_DNHB, default False (0), When 1, launch with grid(N/Br, H, B); when 0, launch with grid(N/Br, B * H).

Key notes

  1. Boolean parsing: ENABLE_FFPA_* booleans and FFPA_PTXAS_VERBOSE are parsed from integer values: 1 = True (enabled), 0 = False (disabled). Exceptions taking string values: FFPA_BUILD_ARCH, FFPA_DEV_HEADDIMS, FFPA_BUILD_STAGES, FFPA_NVCC_DEFINES, FFPA_SM89_PERSIST_KVCFG, ENABLE_FFPA_BUILD_DEBUG (family csv).
  2. Build-time vs. runtime separation: the CUDA extension selection (ENABLE_FFPA_CUDA_IMPL, ENABLE_FFPA_TMA_EXT, ENABLE_FFPA_CUTE_EXT, ENABLE_FFPA_FP8_SM89_EXT, ENABLE_FFPA_QUANT_SPLIT_D, ENABLE_FFPA_F16_ACC, ENABLE_FFPA_CUDA_INPUT_FP16, ENABLE_FFPA_CUDA_MASK_FP32), the stage/headdim/debug sets (FFPA_BUILD_STAGES, FFPA_BUILD_MAX_STAGES, ENABLE_FFPA_ALL_STAGES, ENABLE_FFPA_ALL_HEADDIM, ENABLE_FFPA_BUILD_DEBUG) and the native-kernel tuning knobs below (ENABLE_FFPA_PREFETCH_QKV … ENABLE_FFPA_LAUNCH_GRID_DNHB) are read at build time — changing them requires rebuilding (FFPA_CLEAN=1 bash tools/build_fast.sh). The debug FFPA_* getenv switches (e.g. FFPA_FP8_FORCE_KERNEL) only exist in builds with the matching ENABLE_FFPA_BUILD_DEBUG family. Truly runtime-only knobs: the backend gating, logger, persistent-config and backend-specific groups above.
  3. Default build mode: source builds are Triton-only unless ENABLE_FFPA_CUDA_IMPL=1 is set. FFPA_SKIP_CUDA_EXT=1 is now only an override that force-disables CUDA extension builds even when the CUDA build flags are enabled.
  4. Default build trimming: CUDA source builds are bf16-input-only, carry no fp32-mask variants, no fp16 MMA-acc kernels and no quantized split-D families unless the corresponding flags above are set.
  5. Default arch inference: when FFPA_BUILD_ARCH is unset and ENABLE_FFPA_CUDA_IMPL=1, setup.py queries the currently visible CUDA device's compute capability. On CI / RTD runners with no GPU, either keep the default Triton-only build or explicitly set FFPA_SKIP_CUDA_EXT=1.
  6. MMA Acc combinations: ENABLE_FFPA_FORCE_QK_F16 and ENABLE_FFPA_FORCE_PV_F16 are mutually exclusive options for the FFPA Acc-F32 kernels — enable at most one to choose the mixed-precision split.
  7. Recommended for tuning: enable FFPA_PTXAS_VERBOSE=1 (build-time) and start from the defaults; toggle one runtime knob at a time and measure with python -m ffpa_attn.bench to attribute the speedup or regression.