workload__inference

Task — In-tree inference Workload

Goal

Add one public, in-tree workload named inference that runs eval-only model execution. The recipe chooses the inference mode:

workload: inference
workload_config:
  mode: offline_batch   # offline_batch|continuous_batch

This workload should exercise inference behavior without backward or optimizer steps. It measures latency/throughput, checks outputs for NaN/inf, and returns a normal WorkloadResult.

After this task:

• aorta run --workload inference --steps N runs a public synthetic inference loop.
• Recipes can choose model topology and serving mode.
• One-node smoke works on CPU or GPU.
• Optional single-node multi-rank smoke can run if distributed: true is set in the recipe.

Why Separate From training

Inference has a different failure and performance surface:

• no backward pass
• no optimizer
• eval mode
• prefill/decode latency
• KV cache behavior
• batching behavior
• output/logit stability

Keep training and inference separate workloads so recipes, metrics, and acceptance criteria stay clear.

Recipe vs Sidecar Decision

Use recipes for inference shape:

• inference mode
• model topology
• batch size
• prompt length
• generated token count
• dtype
• KV cache on/off
• deterministic/checksum settings

Use sidecars only for ad-hoc registries: extra mitigations, environments, or temporary local definitions. Do not use sidecars to choose inference mode or model topology.

Files to create / modify

Follow the existing flat-module convention: every in-tree public workload
(workloads/race.py, workloads/llm_determinism.py, workloads/_subprocess.py)
is a single module, not a subpackage. llm_determinism.py already bundles
its config dataclass, helper classes, and the workload in one file — mirror that.
Do not introduce a workloads/inference/ subpackage.

src/aorta/workloads/inference.py                # NEW — InferenceWorkload(Workload)
                                                #       + typed config dataclass(es)
                                                #       + prefill/decode/batching helpers
recipes/example-inference-smoke.yaml            # NEW — offline batch smoke recipe
recipes/example-inference-continuous-smoke.yaml # NEW — continuous batching smoke recipe
tests/workloads/test_inference.py               # NEW — config, lifecycle, result schema

Reuse existing infra — do not write new model builders:

• aorta.models.repeated_block (RepeatedBlockModel + BlockConfig) already
  provides the dense-transformer and MoE (num_experts > 1 → top-1 router)
  topologies. Use it for decoder_transformer / encoder_transformer.
• aorta.data.synthetic_dataset (SyntheticDatasetConfig, create_dataloader)
  provides public synthetic inputs.
• The mlp smoke topology is the only thing that may need a small local builder
  if repeated_block is too heavy for the smallest lifecycle smoke.

Modify pyproject.toml:

[project.entry-points."aorta.workloads"]
inference = "aorta.workloads.inference:InferenceWorkload"

Contract details

• Base contract: import Workload and WorkloadResult from aorta.workloads.
• Name declaration: set name: ClassVar[str] = "inference" to match the existing convention (RaceWorkload.name, LlmDeterminismWorkload.name).
• Launch declaration: default to launch_mode = "single_process" and min_world_size = 1.
• Optional distributed smoke: support a recipe flag like distributed: true later if multi-rank inference is needed. Do not make distributed launch mandatory for the first public inference workload.
• Device behavior: use CUDA/ROCm when available. CPU is acceptable for lifecycle/unit smoke, but result metrics must report the actual device.
• --steps flow: --steps S from aorta run arrives as self.config["steps"] and overrides recipe/default request batches or decode iterations.
• Rank/output behavior: single-process writes through dispatcher as usual. If distributed inference is added later, rank-0-only JSON remains the dispatcher contract.

Recipe shape

schema_version: 1
ticket: EXAMPLE-INFERENCE-SMOKE
workload: inference
trials: 1
steps: 4

workload_config:
  mode: offline_batch        # offline_batch|continuous_batch
  seed: 1234
  device: auto              # auto|cuda|cpu
  dtype: bfloat16           # bfloat16|float16|float32

  request:
    batch_size: 4
    prompt_len: 128
    generate_tokens: 32

  model:
    kind: decoder_transformer # mlp|encoder_transformer|decoder_transformer
    hidden_size: 512
    num_layers: 4
    num_heads: 8
    ffn_size: 2048
    vocab_size: 32000

  serving:
    kv_cache: true   # simulated only — RepeatedBlockModel has no real paged KV cache;
                     # the workload re-uses past hidden states via tensor slicing to
                     # exercise the prefill/decode timing split, not paged attention.
    continuous_batch:
      enabled: false
      max_active_requests: 8
      arrival_pattern: fixed

  checks:
    fail_on_nan_logits: true
    fail_on_nonfinite_output: true
    compare_logits_checksum: true

cells:
  - name: baseline-local
    mitigations: [none]
    environment: local

Implementation guidance

Use typed config dataclasses for structured config. Avoid loose dict plumbing beyond the Workload(config) boundary.

Build topologies from existing in-tree models (aorta.models.repeated_block) and synthetic inputs (aorta.data.synthetic_dataset) — do not author fresh model code:

• mlp: smallest lifecycle smoke (small local builder acceptable if repeated_block is too heavy).
• encoder_transformer: fixed tensor input, useful for embedding/classification-style inference. Map onto RepeatedBlockModel.
• decoder_transformer: synthetic token prompts, prefill + decode loop, optional KV cache. Map onto RepeatedBlockModel.

Note: repeated_block's MoE path is top-1 only (num_experts > 1 selects a top-1 router; there is no top_k knob). If a recipe needs MoE inference, drive it via num_experts and do not expose top_k.

Inference loop should:

• set deterministic seed
• build synthetic token or tensor inputs
• put model in eval() mode
• run with torch.no_grad() or torch.inference_mode()
• run warmup iterations outside measured timings
• measure prefill latency where applicable
• measure decode latency where applicable
• compute throughput metrics
• check logits/outputs for NaN/inf
• optionally compute stable checksums for drift detection
• return WorkloadResult

KV cache note: RepeatedBlockModel has no real paged attention or KV cache. When kv_cache: true, the workload simulates the prefill/decode split by running a full forward pass for prefill then re-running with a shorter input for decode (tensor slicing). This is enough to produce separate prefill and decode latency measurements without implementing real KV caching.

Keep serving modes small:

• offline_batch: fixed batch and fixed token count. This is the required MVP.
• continuous_batch: simple simulated arrivals and active-request scheduling. Keep it deterministic and public-safe.

Do not pull real customer prompts, tokenizers, or model weights into public AORTA.

Acceptance criteria

• inference entry point exists in pyproject.toml.
• InferenceWorkload(Workload) implements setup(), run(), and cleanup().
• InferenceWorkload declares name = "inference", launch_mode = "single_process", and min_world_size = 1.
• Workload lives in a single flat module src/aorta/workloads/inference.py (no subpackage).
• Model topologies reuse aorta.models.repeated_block / aorta.data.synthetic_dataset rather than new model code.
• Recipe config supports mode: offline_batch; continuous_batch can be included if scoped small.
• Recipe config supports at least model.kind: mlp|decoder_transformer.
• Bare smoke works: aorta run --workload inference --trials 1 --steps 2.
• Recipe smoke works: aorta triage run --recipe recipes/example-inference-smoke.yaml --dry-run.
• run() returns WorkloadResult with passed, failure_count, failure_details, step_times_ms, total_iterations, elapsed_sec, main_work_started, executed_iterations, and configured_iterations.
• Metrics include mode, device, dtype, model_kind, parameter_count, batch_size, prompt_len, generate_tokens, prefill_latency_ms, decode_latency_ms, tokens_per_sec, step_time_p50, and step_time_p99 where applicable.
• Numeric checks fail on NaN/inf logits or non-finite outputs.
• Tests cover config parsing, model construction, offline batch run, optional continuous batch scheduling, and WorkloadResult schema.

Out of scope

• Customer/NDA prompts, model weights, traces, or tokenizers.
• Real production serving stack integration.
• Tensor parallelism, pipeline parallelism, or distributed inference.
• RAG with real indexes or private corpora.
• Speculative decoding with real draft/verifier model pairs.
• Throughput tuning.
• search_space() / objective() methods.

How to test

# Unit tests:
python -m pytest tests/workloads/test_inference.py -v

# One-process smoke:
aorta run --workload inference --trials 1 --steps 2

# Recipe dry run:
aorta triage run --recipe recipes/example-inference-smoke.yaml --dry-run

# Recipe smoke:
aorta triage run --recipe recipes/example-inference-smoke.yaml

# Result shape:
python -m json.tool results/inference/trial_d0_m0_t0.json