Environment Variables

Spyre Inference reads the following SPYRE_* environment variables to configure the plugin. Each is evaluated the first time it is read, so set it before launching vLLM. The list below is generated from spyre_inference/envs.py; the comment above each entry documents its effect and default.

environment_variables: dict[str, Callable[[], Any]] = {
    # Comma-separated Spyre device indices to run on, mapping each tensor-parallel
    # local rank to a physical card. Unset lets the runtime pick the default device(s).
    "SPYRE_DEVICES": lambda: os.getenv("SPYRE_DEVICES"),
    # Which model runner the worker uses:
    #  - "v1": the V1 runner (default)
    #  - "v2": the Model Runner V2 port (#906), in bring-up and not yet functional. On
    #    x86 it runs MRV2's Triton kernels on triton-cpu.
    "SPYRE_MODEL_RUNNER": lambda: os.getenv("SPYRE_MODEL_RUNNER") or "v1",
    # Granularity of the decoder's torch.compile graph:
    #  - "block": compile one transformer block at a time (default)
    #  - "model": compile the whole model as a single graph
    "SPYRE_COMPILE_GRANULARITY": lambda: os.getenv("SPYRE_COMPILE_GRANULARITY") or "block",
    # What to do when a model block, attention kernel or the lm_head compiles *after*
    # warmup, which costs a full Inductor compile mid-request:
    #  - "off": nothing
    #  - "warn": log each distinct violation (default)
    #  - "error": raise. Use in CI to keep a warmup-coverage regression from landing,
    #    bearing in mind it catches most but not all: torch runs its compile-start
    #    callbacks only when a process-wide pending counter goes 0 -> 1, so a compile
    #    starting while another is in flight goes unreported.
    # torch-spyre compiles every eager aten op, so those compiles continue for the
    # whole run; they are never reported.
    "SPYRE_COMPILE_GUARD": lambda: os.getenv("SPYRE_COMPILE_GUARD") or "warn",
    # When "1", wrap attention forward/softmax in torch.profiler.record_function
    # spans for kineto trace capture. Off by default: profiled runs are not
    # wall-clock comparable.
    "SPYRE_ATTN_PROFILING": lambda: bool(int(os.getenv("SPYRE_ATTN_PROFILING", "0"))),
    # When "1" (default), pre-compile every attention variant the run can need during
    # warmup, so no request pays an Inductor compile mid-serving. "0" falls back to
    # compiling each variant lazily on first use.
    "SPYRE_ATTN_RECORD": lambda: bool(int(os.getenv("SPYRE_ATTN_RECORD", "1"))),
    # When "1", paged attention walks KV pages and batched decode walks logical
    # block chunks with torch-spyre's `for_each_tile`, so each traced graph holds
    # one loop body. Enabled by default; "0" runs the same bodies under Python loops.
    "SPYRE_ATTN_FOR_EACH_TILE": lambda: bool(int(os.getenv("SPYRE_ATTN_FOR_EACH_TILE", "1"))),
    # Comma-separated kv_len buckets to record, unset uses the SPYRE_ATTN_KV_LADDER preset.
    "SPYRE_ATTN_KV_BUCKETS": lambda: os.getenv("SPYRE_ATTN_KV_BUCKETS"),
    # Named KV bucket ladder: pow2 (default), 8_5, 4_3, 9_8 or uniform.
    "SPYRE_ATTN_KV_LADDER": lambda: os.getenv("SPYRE_ATTN_KV_LADDER") or "pow2",
    # A denser preset for batched decode at SPYRE_ATTN_LARGE_BATCH_MIN_SEQS sequences or
    # more, unioned with SPYRE_ATTN_KV_LADDER. Unset keeps one ladder for every batch size;
    # rejected together with SPYRE_ATTN_KV_BUCKETS.
    "SPYRE_ATTN_KV_LADDER_LARGE_BATCH": lambda: (
        os.getenv("SPYRE_ATTN_KV_LADDER_LARGE_BATCH") or None
    ),
    "SPYRE_ATTN_LARGE_BATCH_MIN_SEQS": lambda: int(
        os.getenv("SPYRE_ATTN_LARGE_BATCH_MIN_SEQS") or "16"
    ),
    # Comma-separated query_len buckets to record, unset uses the default buckets
    # [1] + multiples of min(512, max_num_batched_tokens) up to max_num_batched_tokens.
    "SPYRE_ATTN_QUERY_BUCKETS": lambda: os.getenv("SPYRE_ATTN_QUERY_BUCKETS"),
    # Comma-separated num_seqs buckets for the batched decode kernel, unset uses the
    # default buckets of powers of two from 1 up to max_num_seqs.
    "SPYRE_ATTN_NUM_SEQS_BUCKETS": lambda: os.getenv("SPYRE_ATTN_NUM_SEQS_BUCKETS"),
    # Core cap for the attention compile only, leaving the rest of the model on all 32.
    # "0" (default) lets the LX path pick its own cap and leaves the others uncapped.
    "SPYRE_ATTN_MAX_CORES": lambda: int(os.getenv("SPYRE_ATTN_MAX_CORES", "0")),
    # When "1" (default), enables the batched multi-sequence decode kernel.
    "SPYRE_BATCHED_DECODE": lambda: bool(int(os.getenv("SPYRE_BATCHED_DECODE", "1"))),
    # When "1" (default), a decode batch spanning several KV buckets runs one batched
    # dispatch per bucket group instead of padding every sequence to the longest.
    "SPYRE_JAGGED_DECODE": lambda: bool(int(os.getenv("SPYRE_JAGGED_DECODE", "1"))),
    # When "1" (default), build each linear layer's transposed weight on the host in
    # blocks of source rows. A plain `w.t().contiguous()` copies on one thread; the
    # blocked copy runs on torch's CPU threads and cuts model load time. Bit-identical.
    # Set to "0" when torch runs far more threads than the CPUs it gets (e.g. with
    # SPYRE_UPDATE_THREAD_CONFIG=0 in a CPU-limited pod): there it is slower.
    "SPYRE_BLOCKED_TRANSPOSE": lambda: bool(int(os.getenv("SPYRE_BLOCKED_TRANSPOSE", "1"))),
    # When "1", reuse compiled Spyre kernels across processes by caching them on
    # disk. Off by default. TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 disables the cache
    # even when this flag is enabled.
    "SPYRE_KERNEL_CACHE": lambda: os.getenv("SPYRE_KERNEL_CACHE", "0") == "1",
    # When "1", each TP rank copies only its vocab shard of the logits to the host and the
    # shards are gathered on the CPU. Needs OMP_WAIT_POLICY=PASSIVE (spinning OpenMP
    # threads otherwise slow the next step).
    "SPYRE_LOGITS_CPU_GATHER": lambda: os.getenv("SPYRE_LOGITS_CPU_GATHER", "0") == "1",
    # Maximum number of sequences allowed to prefill in the same batch. "1" (default)
    # serialises prefills, so a batch spends the whole token budget on one prompt
    # instead of topping itself up with a short chunk of the next. Any non-positive
    # value removes the cap, as does a pooling runner, which never decodes.
    "SPYRE_MAX_NUM_PARTIAL_PREFILLS": lambda: int(os.getenv("SPYRE_MAX_NUM_PARTIAL_PREFILLS", "1")),
    # Largest packed-token count handled by a compiled loop of single-token gathered kernels.
    # Larger batches use one all-expert kernel because its fixed weight-read cost is amortized.
    # "8" (default) is measured on gemma-4-26B-A4B: gathered wins up to 8 at TP1 and TP4.
    "SPYRE_MOE_GATHERED_MAX_TOKENS": lambda: int(os.getenv("SPYRE_MOE_GATHERED_MAX_TOKENS", "8")),
    # Override the automatically selected hidden-axis chunk count for gathered MoE. Unset
    # derives it from top_k, stick alignment and the core count.
    "SPYRE_MOE_CHUNKS": lambda: (
        int(value) if (value := os.getenv("SPYRE_MOE_CHUNKS")) is not None else None
    ),
    # CPU budget used to size thread pools. "0" (default) auto-detects the budget
    # (cgroup CPU quota, then physical core count).
    "SPYRE_NUM_CPUS": lambda: int(os.getenv("SPYRE_NUM_CPUS", "0")),
    # When "1" (default), clamp the CPU threading env vars (OMP_NUM_THREADS and
    # friends) to the detected budget to avoid thread oversubscription in
    # CPU-limited containers. Set to "0" to leave them untouched and only warn.
    "SPYRE_UPDATE_THREAD_CONFIG": lambda: bool(int(os.getenv("SPYRE_UPDATE_THREAD_CONFIG", "1"))),
}