Skip to content

Per-Task Ring Sizing

PyPTO exposes Simpler's per-task ring sizing as three optional overrides on RunConfig. They let you size the runtime's per-task ring resources for a single dispatch without touching the compiled artifact or any global state. This works on both the L2 single-chip path (run() / ChipWorker.run()) and the L3 distributed path (DistributedWorker.run() / the one-shot compiled(...)).

The runtime keeps its task-launch resources in ring buffers. Each override maps 1:1 to a field on Simpler's CallConfig.runtime_env and is sized per task submission — per run() / rt.run() call — so different submissions of the same kernel can use different ring sizes. On the L3 path the override is applied per dispatch on top of the program's DistributedConfig baseline (aicpu_thread_num) and reaches every chip in that dispatch.

Field matrix

RunConfig field CallConfig.runtime_env member Controls Constraint
ring_task_window: int \| None ring_task_window Number of in-flight task slots in the task ring power of 2, >= 4
ring_heap: int \| None ring_heap Bytes of the per-ring task-output heap power of 2, >= 1024
ring_dep_pool: int \| None ring_dep_pool Dependency-edge pool capacity [4, INT32_MAX]

None (the default) leaves the field unset (0 on CallConfig). PyPTO writes the value into CallConfig.runtime_env only when it is not None, so an unset field defers entirely to the runtime.

Precedence

For each value the runtime resolves the effective size as:

per-task CallConfig.runtime_env value   (RunConfig override — highest priority)
  └─ falls back to → PTO2_RING_* env var (process-wide)
       └─ falls back to → compile-time default (lowest priority)

So a RunConfig override wins over the PTO2_RING_TASK_WINDOW / PTO2_RING_HEAP / PTO2_RING_DEP_POOL environment variables, which in turn win over the runtime's built-in defaults.

Validation

RunConfig validates the overrides at construction (in __post_init__), mirroring the runtime's RuntimeEnv::validate(). This surfaces a clear ValueError at the call site rather than a deep failure inside the runtime's own CallConfig::validate() at dispatch time:

RunConfig(platform="a2a3", ring_heap=1000)
# ValueError: ring_heap must be a power of 2 >= 1024 (bytes per ring), got 1000

Non-integers (including bool) are rejected with the same ValueError rather than raising an opaque TypeError from the power-of-two check.

Usage

L2 single-chip (run)

from pypto.runtime import run, RunConfig

compiled = run(
    MyProgram,
    a, b, c,
    config=RunConfig(
        platform="a2a3",
        ring_task_window=128,        # 128 in-flight task slots
        ring_heap=8 * 1024 * 1024,   # 8 MiB output heap per ring
        ring_dep_pool=256,           # 256 dependency-edge entries
    ),
)

L3 distributed (DistributedWorker / compiled(...))

Pass the same RunConfig to a per-dispatch call. The prepared worker's shared config is never mutated, so consecutive dispatches — and, in multi-program serving, different programs (e.g. prefill vs. decode) — can each use their own ring sizes:

from pypto.runtime import RunConfig

prefill_cfg = RunConfig(platform="a2a3", ring_task_window=256)
decode_cfg = RunConfig(platform="a2a3", ring_task_window=64)

with compiled.prepare(prefill_cfg) as rt:              # setup once + prewarm
    rt(host_x, weight, host_out, config=prefill_cfg)   # large window for prefill
    rt(host_x, weight, host_out, config=decode_cfg)    # smaller window for decode

# One-shot path honors the same override:
compiled(a, b, c, config=RunConfig(platform="a2a3", ring_heap=8 * 1024 * 1024))

On the L3 dispatch path only the ring_* fields of RunConfig are consumed; the compile-side and DFX fields are ignored (they are set at compile / prepare time, not per dispatch).

Set only the fields you want to override; the rest stay unset and fall back per the precedence above.

Arena prewarm and the single-slot cache

A ring sizing's runtime arena costs ~800 ms to build the first time it is used. Workers build it eagerly at initprepare(config) / ChipWorker / execute_on_device — so the cold build lands in setup instead of inside the first (usually timed) dispatch.

The arena cache is keyed on the full per-ring sizing vector (all four rings' ring_task_window / ring_heap / ring_dep_pool, hashed together). So a single dispatch that sizes rings differently — the list form, e.g. ring_task_window=[256, 128, 64, 0] — is one key and one arena: prewarm covers it completely, as long as prepare(...) gets that same RunConfig. The single-slot limit below is only about different dispatches using different sizing vectors, never about per-ring differences within one dispatch.

The cache holds exactly one sizing vector per worker, and every arena build overwrites that slot. So:

  • Pass the dispatch RunConfig to prepare(...). It is used for the prewarm only — it is not stored, so every dispatch still needs its own config=. In the example above, pass prefill_cfg to both prepare(...) and the first prefill dispatch; the decode dispatch uses decode_cfg.
  • The prewarm only removes the cold build from the first dispatch, and only when that dispatch's sizing matches the prewarmed one — so prewarm the sizing your first dispatch uses, not the most frequent one. Any first dispatch with a different sizing misses, rebuilds, and overwrites the prewarmed slot, wasting the prewarm.
  • Keep the ring sizing constant across a worker's dispatches. The single slot is per worker and is overwritten on every switch, so alternating sizings rebuild the arena on each switch regardless of prewarm (the decode_cfg dispatch above pays a build, and the next prefill_cfg pays another). Splitting the sizings across separate workers is usually not an option — the reason to hold one worker (prepare(..., extra_compiled=[...])) is to share device- resident state such as a KV cache, which separate workers cannot. So pick one sizing that serves all of that worker's dispatch shapes rather than varying it per dispatch.
  • L2 dispatch takes its ring sizing from the per-call RunConfig, so a ChipWorker prewarms the runtime's default sizing; a per-call RunConfig that sizes the rings differently rebuilds once.
  • Simpler's runtime-side implementation: the tensormap_and_ringbuffer runtime under runtime/src/<arch>/runtime/tensormap_and_ringbuffer/ and the RuntimeEnv / CallConfig definitions in runtime/src/common/task_interface/call_config.h.
  • Worker-API examples: runtime/examples/workers/{l2,l3}/per_task_runtime_env/.
  • L3 dispatch entry points that accept the per-dispatch RunConfig: DistributedWorker.run / __call__ and execute_distributed in distributed_runner.py.
  • Orthogonal runtime diagnostics on the same RunConfig: 03-runtime-dfx.md.