Skip to main content

Coding conventions

A short checklist. Skim before opening a PR. None of these are arbitrary; each has bitten the project at least once.

Python style​

  • 4-space indentation, snake_case for functions/variables, PascalCase for classes. Match surrounding code in the file you're editing.
  • No enforced formatter. Don't run black / ruff format on a whole file unless you're rewriting it. Style noise hides real diffs.
  • Imports: keep minimal and consistent. serving/ modules use relative imports (from .scheduler import …).
  • English only in code, comments, log messages, and docstrings. Korean / other-language identifiers and comments will be flagged in review.
  • Docstrings: optional. If you write one, make it a single line that explains why the function exists, not what it does. The signature already says what.
  • No top-level prints. Use serving/core/logger.py (already imported as logger in most files): 
    logger.info(...)
    logger.warning(...)
    logger.success(...)   # Rich-styled green check

CLI flag conventions​

  • CLI flags use hyphens: --cluster-config, --max-num-seqs, --enable-prefix-caching.
  • Internal Python uses underscores: cluster_config, max_num_seqs, enable_prefix_caching.
  • Boolean flags use BooleanOptionalAction so both --enable-X and --no-enable-X work: 
    parser.add_argument('--enable-prefix-caching',
                        action=argparse.BooleanOptionalAction,
                        default=True)
  • Match vLLM naming where applicable (--max-num-batched-tokens, --block-size, --kv-cache-dtype). Users coming from vLLM should not have to relearn.

File and config naming​

  • JSON config filenames: descriptive snake_case (single_node_pim_instance.json, not singleNodePimInstance.json).
  • One config = one scenario. Don't reuse the same cluster JSON across unrelated examples; copy it.
  • Don't commit machine-specific paths. All paths in code and configs must be relative to the repo root.

Things to never do​

These each correspond to a real incident or strong project preference:

  1. Don't add getattr(request, 'attr', default) fallbacks for Request attributes. Initialize all attributes in Request.__init__ and access directly. Fallbacks hide initialization bugs.

  2. Don't assume hidden_size == num_heads * head_dim. Some models (Qwen3) violate this. Always:

    head_dim = config.get('head_dim', n_embd // n_head)
    q_dim   = n_head * head_dim         # NOT n_embd
    kv_dim  = kv_head * head_dim        # NOT n_embd // group

  3. Don't invent layer names. Every name the simulator emits must also appear in the architecture YAML's catalog. Canonical set: qkv_proj, o_proj, gate_up_proj, act_fn, down_proj, rotary_emb, qk_norm, attention, layernorm, final_layernorm, embedding, lm_head, sampler, moe.

  4. Don't edit astra-sim/ unless the change targets simulator integration (Chakra converter, Workload.cc, input configs). Most contributions never touch this directory.

  5. Don't manually edit astra-sim/inputs/*.json. Those files are regenerated by config_builder.py on every run; your edits will be silently overwritten.

  6. Don't commit large generated files. Trace files, outputs/*.csv from your local runs, .et protobufs, and profiler bundle CSVs that exceed the gitignore patterns should stay local. The gitignore is set up; just don't git add -A.

  7. Don't rely on automation to catch you. There are no pre-commit hooks and no test CI: the only GitHub workflow is deploy-docs.yml, which builds the docs site on pushes to main that touch docs/**. Every check in Validating your changes is manual.

  8. Don't add error handling for cases that can't happen. Trust internal invariants; only validate at the boundaries (CLI args, JSON config load, dataset parsing). Defensive programming inside scheduler.py makes the file unreadable.

  9. Don't add features beyond the task at hand. A bug fix doesn't need surrounding cleanup. Three similar lines is better than a premature abstraction.

  10. Don't add comments explaining what the code does. The identifier names already do that. Comments are reserved for why something non-obvious is the way it is (a hidden invariant, a bug workaround, a citation to a paper).

Layer-name and unit reminders​

These two trip up new contributors most often:

  • Profiler CSVs store microseconds (time_us column). The simulator multiplies by 1000 and rounds to nanoseconds at load time. Don't divide twice.
  • Communication sizes for ASTRA-Sim are total (not per-NPU) bytes. ASTRA-Sim divides by ring size internally. If you pass per-NPU sizes, every collective will be N times too small.

Scheduler invariants​

These are the ones with a history. Each has been broken before, and each cost real debugging time.

  • There is no prefill phase and no decode phase. A request just catches up to num_tokens_reached, so num_new = num_tokens_reached - num_computed_tokens — 1 in steady-state decode, the whole remainder for a resumed request. Request.is_prefill() was removed on purpose: it read original_input and so misread a resumed request's recomputation as decoding. Classify for the trace by scheduled token count (> 1 is a prefill chunk, == 1 is decode), never by a phase flag.
  • Never derive sequence length from num_computed_tokens. Preemption resets it to 0. num_tokens_reached is the independent counter, mirroring vLLM's len(_all_token_ids).
  • num_computed_tokens advances at schedule time, as in vLLM's _update_after_schedule, and Batch.scheduled_tokens is the snapshot add_done works from. Advancing it at completion instead lets pp_size > 1 schedule the same tokens twice.
  • Don't add a "preserve the decode state on preemption" special case. num_computed_tokens = 0 is vLLM's own behaviour and is not re-prefill: free_blocks keeps the blocks' hashes, so on re-admission get_computed_blocks finds whatever is still resident and a lower tier returns what was written down. Only the remainder is recomputed. Two earlier attempts to special-case this cost 375 preemptions / 293k recomputed tokens, and 41,569 preemptions / 6 TB of swap.
  • Phase B never preempts, and is skipped entirely on any step that preempted. That anti-thrash rule is load-bearing; without it the running set oscillates preempt → refill → preempt.
  • There is one schedule(), for prefix caching on and off. The pool handles enable_caching=False the way vLLM does — allocate through the same free list, never index — so do not add a second scheduler.

Trace-format invariants​

If you touch trace_generator.py or graph_generator.py:

  • The first layer's input_loc and the last layer's output_loc must be REMOTE:{node_id}. The Chakra converter emits a MEM_LOAD from the first and a MEM_STORE from the last; if either is LOCAL without local memory configured, ASTRA-Sim crashes.
  • The sampler's output_loc and output_size are what feed the MEM_STORE. Don't put them on lm_head: what goes back to the host is the sampled token ids, not the logits.
  • A layer's output_size is not the next layer's input_size, and is not meant to be. qkv_proj emits Q+K+V while rotary_emb declares only Q+K, and attention reads K/V from the cache rather than the activation. There is no chain to "restore".
  • Never split pipeline stages by trace-line count. A stage may only be cut on a transformer-block boundary, the one place where the upstream output_size and downstream input_size are the same tensor (the hidden state). ASTRA-Sim's analytical backend keys its send/recv tracker on (tag, src, dst, chunk_size, chunk_id), so a size mismatch silently deadlocks the receiving NPU instead of raising. That is what pp_stage_boundaries in the trace header exists for.
  • Don't "restore" log-space interpolation in _axis_bracket because the profiler's sweep grid is geometric. Grid spacing decides where the kernel is sampled; the blend decides how two samples combine; the kernel is linear in each axis. Log blending biased estimates 11.6-14.4% high against 2.3-3.7% for linear, measured leave-one-out across every bundle in profiler/perf/.

Commit and PR style​

The short version (full process is on PR workflow):

  • Commit messages: short imperative one-liner.
    • Good: Fix incorrect evict_size accumulation, Add Qwen3 model support.
    • Bad: fixes, update scheduler.py, WIP.
  • One logical change per commit. Don't bundle a refactor with a feature.
  • PR description includes the validation command you ran, so the reviewer can rerun it — usually ./serving/validate.sh, plus the report table if anything moved.

What's next​