# ATOM configuration guide

ATOM (AiTer Optimized Model) is AMD's lightweight LLM inference engine built on
[AITER](https://github.com/ROCm/aiter) kernels for ROCm/HIP GPUs. This guide
documents every configuration class, CLI flag, and environment variable that
controls ATOM's runtime behaviour.

## Quick reference

| Config Class | Primary Purpose |
|---|---|
| `Config` | Master dataclass — model path, memory, TP size, scheduler limits, KV cache, profiler, and references to all sub-configs |
| `CompilationConfig` | Compilation level (0-3), CUDA graph capture sizes, piecewise splitting ops, inductor settings |
| `CompilationLevel` | Integer constants for the four compilation levels |
| `CUDAGraphMode` | Enum controlling how CUDA graphs are captured (`NONE` / `TR_HYBRID` token graphs / `FULL` / `FULL_AND_TR_HYBRID`) |
| `QuantizationConfig` | Layer-wise quantization orchestrator: global config, per-layer overrides, exclude lists, layer name remapping |
| `LayerQuantConfig` | Per-layer quantization spec (frozen dataclass): quant type, dtype, dynamic flag, method |
| `ParallelConfig` | Data-parallel size, rank, master IP/port |
| `SpeculativeConfig` | Speculative decoding method, draft model, number of speculative tokens |
| `KVCacheConfig` / `KVCacheTensor` | Per-layer KV cache tensor descriptors (k/v caches and scales) |
| `SamplingParams` | Temperature, max tokens, stop strings, ignore-EOS flag |
| `EngineArgs` | CLI argument parser that builds a `Config` for `LLMEngine` |

## Master configuration (`Config`)

Defined in `atom/config.py`. The root dataclass that the engine consumes.

| Field | Type | Default | Description |
|---|---|---|---|
| `model` | `str` | *(required)* | HuggingFace model name or local path |
| `trust_remote_code` | `bool` | `False` | Trust remote code when loading the model from HuggingFace |
| `max_num_batched_tokens` | `int` | `16384` | Maximum number of tokens batched together per scheduler step |
| `scheduler_delay_factor` | `float` | `0.0` | Multiplicative delay (factor x previous prompt latency) before scheduling the next prompt |
| `max_num_seqs` | `int` | `512` | Maximum number of sequences batched together |
| `max_model_len` | `int \| None` | `None` | Maximum context length; defaults to `hf_config.max_position_embeddings` (capped by it when set) |
| `gpu_memory_utilization` | `float` | `0.9` | Fraction of GPU memory available for KV cache and weights (0.0 — 1.0) |
| `tensor_parallel_size` | `int` | `1` | Number of tensor-parallel GPUs (1 — 8) |
| `enforce_eager` | `bool` | `False` | Disable compilation and CUDA graphs; run in eager mode |
| `parallel_config` | `ParallelConfig` | `ParallelConfig()` | Data-parallel configuration (see Section 4) |
| `kv_cache_block_size` | `int` | `16` | Block size for paged KV cache; must be a multiple of 16 or exactly 1. The FP4 sparse indexer requires exactly 64 |
| `num_kvcache_blocks` | `int` | `-1` | Number of KV cache blocks (`-1` = auto) |
| `kv_cache_dtype` | `str` | `"bf16"` | KV cache data type (`"bf16"` or `"fp8"`) |
| `index_cache_dtype` | `str \| None` | `None` | Indexer-cache dtype, resolved after model detection. Native single-node DeepSeek-V4 defaults to `"fp4"` except on gfx942; plugin and KV-transfer integrations retain `"fp8"`. DeepSeek-V4.1 defaults to `"fp8"`; its runtime also takes `"fp4"` (the official E2M1 + E8M0 arithmetic in the row-group scorer's page-8 layout, values and scales in two planes, which needs `candidate_block_size` 8) and refuses any other format before weights load, rather than overridden here. Other models inherit `kv_cache_dtype`. An explicit `"bf16"`, `"fp8"`, or `"fp4"` value is preserved. `"fp4"` is honoured only where the FP4 mqa-logits kernels apply -- gfx950, an indexer head dim of 128, an indexer head count that is a multiple of 16, and `--block-size 64`; anything else logs one warning and scores in FP8. It runs under DCP, and is refused outright under PCP and KV transfer rather than falling back. |
| `enable_prefix_caching` | `bool` | `False` | Enable prefix caching to reuse KV blocks across requests sharing the same prefix |
| `enable_log_stats` | `bool` | `True` | Emit the periodic engine-status line (running/waiting reqs, KV usage, prefix-cache hit rate, prompt/generation throughput). Applies to offline `LLM(...)` as well as to the server. Scoped to that line: `[MTP Stats]` and `[Cache Stats]` have their own gates |
| `throughput_log_interval` | `float` | `10.0` | Seconds between engine-status lines. Must be > 0 |
| `cache_hit_rate_window` | `int` | `1000` | Requests in the sliding window behind the engine-status line's prefix-cache hit rate. Must be > 0. Only that line is windowed; `/metrics` and `[Cache Stats]` stay cumulative |
| `enable_dynamic_chunking` | `bool` | `False` | Size pipeline-parallel prefill chunks from a calibrated equal-latency model. Needs `pipeline_parallel_size > 1` and chunked prefill, and does not apply to RapidServe or a P/D decode instance (`kv_role: kv_consumer`); otherwise the flag is turned off with a warning and chunking stays fixed |
| `dynamic_chunking_smooth_factor` | `float \| None` | `None` | Blend from the initial chunk (`0`) to the equal-latency chunk (`1`). `None` takes `ATOM_DYNAMIC_CHUNKING_SMOOTH_FACTOR`, read only when the feature is on. Must be in `[0, 1]` when the feature is on |
| `dynamic_chunking_min_chunk_size` | `int` | `4096` | Floor for a solved chunk. Must be positive when the feature is on |
| `state_checkpoint_interval_tokens` | `int` | `8192` | For models with per-request state (DeepSeek-V4 compressor ring, GDN recurrent state, Kimi-K3 KDA): tokens between rungs of the checkpoint ladder. The sign carries three policies — `>0` a rung every N tokens (must be a multiple of the prefix-cache hash block size), `0` checkpointing off entirely, `-1` ladder off while the prompt-end anchor and the demand rung still place. See the state-checkpoint section of the [scheduling & KV cache guide](scheduling_kv_cache_guide.md) |
| `state_checkpoint_demand` | `bool` | `True` | Whether a prefix hit refused for want of a checkpoint may place a rung of its own. Off leaves the prompt-end anchor as the only placement. Overridden by `ATOM_STATE_CHECKPOINT_DEMAND` |
| `enable_decoder_swa_bounded_replay` | `bool` | `True` | DeepSeek-V4.1 only: prefill runs the decoder (layers 20..39) only on the rows the next forward reads, its windows truncated there (the V4.1 report's SWA Bounded Replay). Faster prefill, not numerically equivalent to a full prefill. See [DeepSeek-V4.1 runtime](deepseek_v41_runtime.md#decoder-swa-bounded-replay) |
| `enable_encoder_swa_bounded_replay` | `bool` | `True` | DeepSeek-V4.1 only: a prefix hit resumes without a state checkpoint, rebuilding the sliding-window state from the hash block below the hit; no state checkpoint is kept. See [DeepSeek-V4.1 runtime](deepseek_v41_runtime.md#encoder-swa-bounded-replay) |
| `min_final_prefill_rows` | `int` | `0` | Set by the model, not a flag: the fewest rows the forward finishing a prompt runs; one resuming closer to the end starts earlier, on the hash-block grid (`BlockManager.replay_start`). DeepSeek-V4.1 with decoder bounded replay: `sliding_window` |
| `replay_history_tokens` | `int` | `0` | Set by the model, not a flag: the tokens before a forward that starts below where its prefill resumed which the backend reads (`Replay.history_ids`). DeepSeek-V4.1: its Engram n-gram lookback, `engram_max_ngram_size - 1` |
| `state_replay_tokens` | `int` | `0` | Set by the model, not a flag: how far below a prefix hit with no state checkpoint behind it a prefill starts to rebuild that state exactly (`BlockManager.replay_start`); 0 takes no such hit. DeepSeek-V4 under `ATOM_DSV4_STATE_REPLAY`: the window receptive field `W + (L-1)(W-1)`, 7,748 on V4-Pro and 7,875 with MTP or DSpark |
| `port` | `int` | `8006` | Engine internal communication port |
| `torch_profiler_dir` | `str \| None` | `os.getenv("ATOM_TORCH_PROFILER_DIR", None)` | Directory for saving PyTorch profiler traces; creates the directory if it does not exist |
| `compilation_config` | `CompilationConfig` | `CompilationConfig()` | Compilation and CUDA graph settings (see Section 2) |
| `quant_config` | `QuantizationConfig` | *(auto-detected)* | Quantization settings; auto-detected from HuggingFace config during `__post_init__` via `QuantizationConfig(hf_config)` (see Section 3) |
| `asyncio_mode` | `bool` | `False` | Enable asyncio-based engine loop |
| `load_dummy` | `Optional[str]` | `None` | Dummy-weight mode (no checkpoint read): `None` off; `"empty"` skip load (uninitialized, legacy); `"zero"` all-zero; `"xavier"` xavier for bf16, constant target magnitude for fp4/fp8 |
| `enable_expert_parallel` | `bool` | `False` | Enable Expert Parallelism for MoE models |
| `master_addr` | `str` | `"127.0.0.1"` | Master address for distributed communication |
| `capture_sizes` | `Optional[list[int]]` | `None` | Explicit list of batch sizes for CUDA graph capture; derived from `compilation_config` during init |
| `enable_dp_attention` | `bool` | `False` | Enable data-parallel attention |
| `dp_load_balance` | `str` | `"least_requests"` | DP request-routing strategy: `"round_robin"` (legacy), `"least_requests"` (default; fewest in-flight requests, ties broken by lighter in-flight prompt-token load), or `"least_tokens"` (lowest `sum_prompt_tokens + ATOM_DP_LB_REQ_EQUIV * num_reqs`). Only effective when >1 DP rank. See distributed guide §2 |
| `torch_dtype` | `torch.dtype` | *(computed)* | Inferred from `hf_config.torch_dtype`; falls back to `torch.bfloat16` |
| `speculative_config` | `Optional[SpeculativeConfig]` | `None` | Speculative decoding configuration (see Section 5) |
| `bos_token_id` | `int` | `-1` | Beginning-of-sequence token ID (`-1` = use model default) |
| `eos_token_id` | `int` | `-1` | End-of-sequence token ID (`-1` = use model default) |
| `stop_token_ids` | `list[int]` | `[]` | Additional stop token IDs; populated from `GenerationConfig.eos_token_id` during init |

**Auto-derived fields** (set in `__post_init__` or by `ModelRunner.get_num_blocks()`, not user-supplied):

| Field | Type | Description |
|---|---|---|
| `hf_config` | `PretrainedConfig` | Loaded automatically via `get_hf_config(model)` |
| `generation_config` | `GenerationConfig` | Loaded automatically via `get_generation_config(model)` |
| `pool_entries` | `dict[str, int]` | Entries sized for each cache class the attention builders declared via `AttentionMetadataBuilder.sub_pool_specs()` — the paged KV blocks, plus per-request STATE classes (GDN recurrent state, the DeepSeek-V4 compressor ring and sliding-window pool). Computed by `ModelRunner.get_num_blocks()` in the runner subprocess and carried to the engine process, where each consumer looks up the class it declared |
| `pool_entries_per_req` | `dict[str, int]` | Per-request multiplicity of each class, so a consumer can turn an entry count into a request count (`entries // entries_per_req`); same origin as `pool_entries` |

## Compilation configuration (`CompilationConfig`)

Defined in `atom/config.py`. Controls torch.compile and CUDA graph behaviour.

### Compilation levels (`CompilationLevel`)

| Constant | Value | Description |
|---|---|---|
| `NO_COMPILATION` | `0` | No compilation — pure eager execution |
| `DYNAMO_AS_IS` | `1` | Use torch.compile / TorchDynamo as-is |
| `DYNAMO_ONCE` | `2` | TorchDynamo with a single compilation pass |
| `PIECEWISE` | `3` | Piecewise compilation: the graph is split at `splitting_ops` and each piece compiled by Inductor (recommended for production; CUDA graphs are chosen separately by `CUDAGraphMode`) |

### `CompilationConfig` fields

| Field | Type | Default | Description |
|---|---|---|---|
| `level` | `int` | `0` | Compilation level (see table above); must be 0 — 3 |
| `use_cudagraph` | `bool` | `True` | Whether to use CUDA graphs |
| `cudagraph_capture_sizes` | `Optional[list[int]]` | `None` | Explicit list of batch sizes for CUDA graph capture; overrides `cuda_graph_sizes` when set |
| `cuda_graph_sizes` | `list[int]` | `[]` (post-init: `[512]`) | CUDA graph sizing strategy: 1 value generates `[1..8] + range(16, N+1, 16)`; multiple values used as-is; empty defaults to `[512]` |
| `debug_dump_path` | `str` | `""` | Path to dump debug / compilation information |
| `cache_dir` | `str` | `""` | Directory for compilation caches |
| `use_inductor` | `bool` | `True` | Enable TorchInductor backend |
| `cudagraph_mode` | `Optional[CUDAGraphMode]` | `None` | CUDA graph capture mode (see [CUDA graph mode](#cuda-graph-mode-cudagraphmode)); set to `TR_HYBRID` automatically at level 3 |
| `splitting_ops` | `Optional[list[str]]` | `None` | Ops that split the graph into sub-graphs for piecewise compilation; auto-populated at level 3 with `["aiter.unified_attention_with_output", "aiter.mla_attention"]` |
| `cudagraph_copy_inputs` | `bool` | `False` | Copy input tensors into internally managed buffers before CUDA graph replay; effective at level 3 whenever `cudagraph_mode` is not `NONE` |
| `compile_sizes` | `Optional[list[Union[int, str]]]` | `None` | Sizes to compile for inductor; accepts integers and the string `"cudagraph_capture_sizes"` |
| `inductor_compile_config` | `dict` | `{}` | Additional configuration passed to the inductor backend |

### CUDA graph mode (`CUDAGraphMode`)

| Mode | Value | Description |
|---|---|---|
| `NONE` | `0` | No CUDA graph capture |
| `TR_HYBRID` | `1` | Token graphs and request graphs mixed: a token graph per token bucket, breaking only where a launch depends on host data (default at level 3) |
| `FULL` | `2` | Full CUDA graph capture for all batches; best for small models / short prompts |
| `FULL_DECODE_ONLY` | `(FULL, NONE)` | Full CUDA graphs for decode batches only; mixed prefill-decode runs without graphs (useful in P/D setups) |
| `FULL_AND_TR_HYBRID` | `(FULL, TR_HYBRID)` | Recorded as `TR_HYBRID` by the native runner: decode replays token graphs too, no full decode graph |

Helper methods on `CUDAGraphMode`:

- `decode_mode()` — returns the mode used for pure decode batches.
- `mixed_mode()` — returns the mode used for mixed prefill-decode batches.
- `uses_tr_hybrid()` — whether the mode records token graphs.
- `has_full_cudagraphs()` — whether the mode includes full CUDA graph capture.
- `separate_routine()` — whether decode and mixed batches use different routines.

## Quantization configuration (`QuantizationConfig` & `LayerQuantConfig`)

Defined in `atom/config.py` and `atom/quant_spec.py`. The quantization system uses two classes:

- **`QuantizationConfig`** — the top-level orchestrator that holds a global config, per-layer overrides, and exclusion lists.
- **`LayerQuantConfig`** — a frozen dataclass (defined in `atom/quant_spec.py`) that stores the concrete quantization parameters for a single layer or as the global default. Typed, immutable, with attribute access (e.g., `spec.quant_type`).

### `LayerQuantConfig` fields

`LayerQuantConfig` is a frozen dataclass. Fields are accessed as typed attributes (e.g., `spec.quant_type`).

| Field | Type | Default | Description |
|---|---|---|---|
| `quant_type` | `QuantType` | `QuantType.No` | Quantization granularity (see [`QuantType` values](#quanttype-values-from-aiter)) |
| `quant_dtype` | `torch.dtype` | `torch.bfloat16` | Data type for quantized weights |
| `is_dynamic` | `bool` | `True` | Use dynamic quantization (scales computed at runtime) |
| `quant_method` | `str` | `""` | Quantization method (e.g., `"quark"`, `"compressed-tensors"`) |

### `QuantizationConfig` attributes

| Attribute | Type | Description |
|---|---|---|
| `torch_dtype` | `torch.dtype` | The model's default dtype (from `hf_config.torch_dtype`) |
| `hf_quant_config` | `dict \| None` | Raw `quantization_config` dict from HuggingFace config |
| `global_quant_config` | `LayerQuantConfig` | Default quantization spec applied to all layers |
| `_parsed.layer_pattern_specs` | `list[tuple[str, LayerQuantConfig]]` | Per-layer overrides keyed by layer name pattern (supports fnmatch globs like `"*.mlp.*"`) |
| `exclude_layers` | `list[str]` | Layer names excluded from quantization (supports exact match and `"re:"` regex prefix) |
| `quant_method` | `str` | Top-level quantization method name (e.g., `"quark"`, `"compressed-tensors"`) |

Key methods:

| Method | Description |
|---|---|
| `get_name()` | Returns the quantization method name |
| `get_layer_quant_config(layer_name)` | Returns the `LayerQuantConfig` for a layer: checks exclusions first, then per-layer overrides, then falls back to global spec |
| `remap_layer_name(hf_config, packed_modules_mapping)` | Remaps layer names for packed/fused modules (e.g., `q_a_proj` → `fused_qkv_a_proj` for DeepSeek) |
| `compute_hash()` | Returns a SHA-256 hash of the quantization config for cache invalidation |


### `QuantType` values (from AITER)

| Value | Description |
|---|---|
| `QuantType.No` | No quantization |
| `QuantType.per_Token` | Per-token / per-channel quantization |
| `QuantType.per_1x128` | Block quantization with group size 128 |
| `QuantType.per_1x32` | Block quantization with group size 32 |
| `QuantType.per_128x128` | Large 2D block quantization (remapped to `per_1x128` in MoE kernels) |
| `QuantType.per_Tensor` | Per-tensor quantization |

### Supported quantization dtypes

| Dtype | AITER Key | Notes |
|---|---|---|
| FP8 (E4M3) | `"fp8"` | 8-bit floating point |
| MXFP4 | `"fp4x2"` | Microscaling FP4; forces `QuantType.per_1x32` |
| INT8 | `"i8"` | 8-bit integer |
| INT4 | `"i4x2"` | 4-bit integer (packed) |

### Auto-detection from HuggingFace

During `Config.__post_init__`, ATOM constructs `QuantizationConfig(hf_config)` which
reads `hf_config.quantization_config` and automatically determines quantization
parameters:

**For quark models** (`quant_method == "quark"`):

1. Parses `global_quant_config` dict via `QuarkParser` to produce the global `LayerQuantConfig`.
2. Parses each entry in `layer_quant_config` dict to produce per-layer overrides.
3. Reads the `"exclude"` list for excluded layers.
4. Within each config dict, `weight.qscheme` determines `quant_type` (`"per_channel"` → `per_Token`, `"per_tensor"` → `per_Tensor`, `"per_group"` → `per_1x32`), and `weight.dtype` determines `quant_dtype`.
5. `input_tensors.is_dynamic` controls dynamic quantization (defaults to `True` if absent).

**For other models** (compressed-tensors, etc.):

1. If `quant_method == "compressed-tensors"` or channel quantization is detected, sets `per_Token`.
2. If `weight_block_size` or `group_size` is found: group size 128 maps to `per_1x128`, group size 32 maps to `per_1x32`.
3. Otherwise falls back to `per_Tensor`.
4. The dtype is parsed from fields like `dtype`, `weight_dtype`, or `quant_method` looking for `fp8`, `fp4`, `mxfp4`, `int8`, `int4`, or `num_bits`.
5. If `activation_scheme` is `"static"`, `is_dynamic` is set to `False`.
6. Excluded layers are read from the `"ignore"` key.

### Layer-level quantization dispatch

Linear layers, MoE layers, and fused ops call `quant_config.get_layer_quant_config(prefix)` to obtain the appropriate `LayerQuantConfig` for their position in the model. This enables mixed-precision quantization where different layers can have different quant types and dtypes (e.g., FP8 for attention, FP4 for MLP).

### Online quantization at load time

ATOM can re-quantize model weights while loading them by passing
`--online_quant_config` to the engine. This is useful when the source checkpoint
is unquantized or uses a different supported quantization layout than the runtime
layout you want to benchmark.

> For ready-to-run recipes (DeepSeek-R1, Qwen3 MoE, …), pattern-design guidance,
> verification, and troubleshooting, see
> [`online_quantization_guide.md`](./online_quantization_guide.md). This section
> documents only the field-level reference.

The flag accepts a JSON object:

```bash
--online_quant_config '{
  "global_quant_config": "ptpc_fp8",
  "layer_quant_config": {"*expert*": "mxfp4"},
  "exclude_layer": ["lm_head", "*.gate.*"]
}'
```

Fields:

| Field | Type | Description |
|---|---|---|
| `global_quant_config` | `str` | Default target quantization format for all layers. |
| `layer_quant_config` | `dict[str, str]` | Per-layer target overrides. Keys are fnmatch-style layer patterns such as `"*expert*"`. |
| `exclude_layer` | `str \| list[str]` | Layers to leave unquantized. Supports exact/prefix matches, glob patterns, and `re:` regex entries. Prefer a JSON list when excluding multiple patterns. |

Supported target format strings:

| Format | Target config |
|---|---|
| `ptpc_fp8` | `QuantType.per_Token` with FP8 weights |
| `mxfp4` | `QuantType.per_1x32` with packed MXFP4 weights |

Notes:

- An empty JSON object (`--online_quant_config '{}'`) is treated the same as not passing the flag and does not enable online quantization.
- Online quantization currently applies when the source model is unquantized or uses supported FP8 block quantization (`QuantType.per_1x128`). The loader dequantizes FP8 block weights before applying the requested target format.
- Tensor-parallel weights are gathered before quantization only when local quantization would produce different scales than quantizing the full unpartitioned weight.
- Rank 0 writes an `online_quant_info_*.json` summary with elapsed time and per-layer target formats. The file is written under `ATOM_TORCH_PROFILER_DIR` when set, otherwise the current working directory.

## Parallel configuration (`ParallelConfig`)

Defined in `atom/config.py`. Controls data parallelism. Environment variables
(Section 8) override defaults when set.

| Field | Type | Default | Description |
|---|---|---|---|
| `data_parallel_size` | `int` | `1` | Number of data-parallel groups; overridden by `ATOM_DP_SIZE` env var |
| `data_parallel_size_local` | `int \| None` | `None` → `data_parallel_size` | DP ranks on **this** node. Defaults to the global size, i.e. single-node. Set it lower to give a node one slice of a multi-node run; also reaches MoRI as `gpu_per_node`. Overridden by `ATOM_DP_SIZE_LOCAL` |
| `data_parallel_rank` | `int` | `0` | First **global** DP rank owned by this node; overridden by `ATOM_DP_RANK` |
| `data_parallel_rank_local` | `Optional[int]` | `None` | Local rank within the data-parallel group (SPMD mode); overridden by `ATOM_DP_RANK_LOCAL` |
| `data_parallel_master_port` | `int` | `29500` | Port used by the data-parallel master for process group initialization |
| `data_parallel_base_port` | `int` | `0` | Model-runner TCPStore port; automatically bound for single-node runs, explicitly shared across nodes |
| `data_parallel_master_ip` | `str` | `"127.0.0.1"` | IP address of the data-parallel master |

**Computed property:**

- `world_size` — set during init, equals TP x PP.
- `world_size_across_dp` — `world_size * data_parallel_size`.

## Speculative decoding configuration (`SpeculativeConfig`)

Defined in `atom/config.py`. Currently only the Multi-Token Prediction (MTP)
method with `num_speculative_tokens=1` is supported.

| Field | Type | Default | Description |
|---|---|---|---|
| `method` | `Optional[str]` | `""` | Speculative decoding method; currently only `"mtp"` is accepted |
| `model` | `Optional[str]` | `None` | Draft model name or path (typically the same as the target model for MTP) |
| `num_speculative_tokens` | `Optional[int]` | `None` | Number of speculative tokens per iteration; **must be `1`** |
| `draft_model_hf_config` | `Optional[PretrainedConfig]` | `None` | HuggingFace config for the draft model; auto-loaded from `model` when `None` |

### Table-driven MTP config

MTP configuration uses two class-level lookup tables to support multiple model
families without per-model branching.

**`_MTP_TYPE_MAP`** — maps a base `model_type` to its MTP `model_type`:

| Base `model_type` | MTP `model_type` |
|---|---|
| `deepseek_v3` | `deepseek_mtp` |
| `glm_moe_dsa` | `deepseek_mtp` |
| `qwen3_next` | `qwen3_next_mtp` |
| `qwen3_5` | `qwen3_5_mtp` |
| `qwen3_5_moe` | `qwen3_5_mtp` |
| `qwen3_5_text` | `qwen3_5_mtp` |
| `qwen3_5_moe_text` | `qwen3_5_mtp` |

**`_MTP_CONFIG`** — maps MTP `model_type` to a `(n_predict_attr, architecture)` tuple:

| MTP `model_type` | `n_predict_attr` | Architecture |
|---|---|---|
| `deepseek_mtp` | `num_nextn_predict_layers` | `DeepSeekMTPModel` |
| `qwen3_next_mtp` | `num_nextn_predict_layers` | `Qwen3NextMTPModel` |
| `qwen3_5_mtp` | `mtp_num_hidden_layers` | `Qwen3_5MTPModel` |

### Post-init behaviour (`hf_config_override`)

The static method `hf_config_override` applies a two-step transformation to the
draft model's HuggingFace config:

1. **Resolve model type** — looks up `hf_config.model_type` in `_MTP_TYPE_MAP`.
   If found, rewrites `model_type` to the MTP variant (e.g.
   `deepseek_v3` -> `deepseek_mtp`).

2. **Apply MTP overrides** — looks up the (possibly rewritten) `model_type` in
   `_MTP_CONFIG`. If found:
   - Reads `n_predict` from the model-specific attribute (e.g.
     `num_nextn_predict_layers` or `mtp_num_hidden_layers`), defaulting to 1.
     Warns and forces it to 1 if the original value differs.
   - Sets `n_predict=1`, `num_nextn_predict_layers=1` (universal across all MTP
     families), and `architectures` to the corresponding MTP model class.
   - **Qwen3.5 MTP only**: additionally injects `n_shared_experts=1` and
     `n_routed_experts` (read from `hf_config.num_experts`, default 0) so the
     MTP module can construct its MoE layer.

Other post-init steps:

- Loads `draft_model_hf_config` from `model` if not provided.
- Extracts `text_config` from multimodal model configs when present.
- `Config.__post_init__` raises `ValueError` if `num_speculative_tokens != 1`.

## Sampling parameters (`SamplingParams`)

Defined in `atom/sampling_params.py`. Passed per-request to control generation.

| Field | Type | Default | Description |
|---|---|---|---|
| `temperature` | `float` | `1.0` | Sampling temperature; lower values make output more deterministic |
| `max_tokens` | `int` | `64` | Maximum number of tokens to generate |
| `ignore_eos` | `bool` | `False` | Continue generating past the EOS token |
| `stop_strings` | `Optional[list[str]]` | `None` | List of strings that trigger generation to stop |

## CLI arguments (`EngineArgs`)

Defined in `atom/model_engine/arg_utils.py`. The `EngineArgs` dataclass exposes
all flags via `add_cli_args()` and converts them into a `Config` via
`create_engine()`.

| Flag | Short | Type | Default | Description |
|---|---|---|---|---|
| `--model` | | `str` | `"Qwen/Qwen3-0.6B"` | Model name or path |
| `--trust-remote-code` | | flag | `False` | Trust remote code when loading model |
| `--tensor-parallel-size` | `-tp` | `int` | `1` | Tensor parallel size |
| `--data-parallel-size` | `-dp` | `int` | `1` | Data parallel size |
| `--enforce-eager` | | flag | `False` | Enforce eager mode execution |
| `--enable_prefix_caching` | | flag | `False` | Enable prefix caching |
| `--enable-log-stats` / `--no-enable-log-stats` | | flag | `True` | Emit the periodic engine-status line |
| `--throughput-log-interval` | | `float` | `10.0` | Seconds between engine-status lines |
| `--port` | | `int` | `8006` | Engine internal port |
| `--kv_cache_dtype` | | `str` | `"bf16"` | KV cache dtype; choices: `bf16`, `fp8` |
| `--index-cache-dtype`, `--index_cache_dtype` | | `str` | `None` | Indexer-cache dtype; choices: `bf16`, `fp8`, `fp4`. When omitted, uses the architecture- and integration-aware `Config.index_cache_dtype` defaults described above. |
| `--block-size` | | `int` | `16` | KV cache block size (maps to `kv_cache_block_size`); the FP4 sparse indexer requires exactly 64 |
| `--max-model-len` | | `int` | `None` | Maximum model context length; defaults to `hf_config.max_position_embeddings` |
| `--cudagraph-capture-sizes` | | `str` | `"[1,2,3,4,5,6,7,8,16,32,48,64,128,256,512]"` (every batch of 1..8: a padded one runs the next size's step) | CUDA graph capture sizes as a Python list string |
| `--level` | | `int` | `3` | Compilation level (0 — 3) |
| `--load_dummy` | | `{empty,zero,xavier}` (optional value) | `None` | Dummy weights: bare/`=empty` skip load; `=zero` all-zero; `=xavier` xavier(bf16)/constant-magnitude(fp4/fp8) |
| `--enable-expert-parallel` | | flag | `False` | Enable Expert Parallelism (EP MoE) |
| `--torch-profiler-dir` | | `str` | `None` | Directory for torch profiler traces |
| `--enable-dp-attention` | | flag | `False` | Enable DP attention |
| `--method` | | `str` | `None` | Speculative method; choices: `mtp` |
| `--num-speculative-tokens` | | `int` | `1` | Number of speculative tokens per iteration |
| `--max-num-batched-tokens` | | `int` | `16384` | Maximum number of tokens to batch in the async engine |
| `--state-checkpoint-interval-tokens` | | `int` | `8192` | Tokens between rungs of the per-request state checkpoint ladder; must be a multiple of the prefix-cache hash block size. `0` turns checkpointing off entirely; **`-1` turns off only the ladder**, leaving the prompt-end anchor and the demand rung as the placements. Prompts shorter than one interval publish no rung. A rung also quantizes prefill chunk boundaries, since a checkpoint is only valid where a forward ends exactly on one — which is most of what `-1` buys back |
| `--state-checkpoint-demand` / `--no-state-checkpoint-demand` | | `bool` | on | Let a prefix hit that was refused for want of a checkpoint place a rung of its own. Off leaves the prompt-end anchor as the only placement. `ATOM_STATE_CHECKPOINT_DEMAND=0` overrides this without touching a launch script |
| `--decoder-swa-bounded-replay` / `--no-decoder-swa-bounded-replay` | `--enable-decoder-swa-bounded-replay` | `bool` | on | DeepSeek-V4.1 only: the forward finishing a prompt runs the decoder (layers 20..39) on its last 128 rows only, their windows truncated at the first of them, and layer 39's query on the last row; one that does not finish runs none of it past layer 20's compression (the V4.1 report's SWA Bounded Replay). Not numerically equivalent to a full prefill. No effect on other models |
| `--encoder-swa-bounded-replay` / `--no-encoder-swa-bounded-replay` | `--enable-encoder-swa-bounded-replay` | `bool` | on | DeepSeek-V4.1 only: a prefix hit resumes without a state checkpoint, recomputing the hash block below the hit to rebuild the sliding-window state (no cached row is rewritten); no state checkpoint is kept, so none cuts a prefill chunk. Not numerically equivalent to resuming a checkpoint. No effect on other models |
| `--max-num-seqs` | | `int` | `512` | Maximum number of sequences to batch together |
| `--gpu-memory-utilization` | | `float` | `0.9` | Fraction of GPU memory to use (0.0 — 1.0) |
| `--scheduler-delay-factor` | | `float` | `0.0` | Delay factor multiplied by previous prompt latency before scheduling next prompt |
| `--online_quant_config` | | JSON string | `None` | Load-time online quantization override; see Section 3.7 |

**Example:**

```bash
python -m atom.entrypoint \
    --model deepseek-ai/DeepSeek-R1 \
    --tensor-parallel-size 8 \
    --level 3 \
    --cudagraph-capture-sizes "[1,2,4,8,16,32,64,128,256]" \
    --kv_cache_dtype fp8 \
    --gpu-memory-utilization 0.92 \
    --max-num-seqs 256
```

## Environment variables

### Variables registered in `atom/utils/envs.py`

All variables use lazy evaluation. Boolean variables treat `"1"` as `True` and
anything else (including unset) as `False`, unless noted otherwise.

| Variable | Type | Default | Description |
|---|---|---|---|
| `ATOM_DP_RANK` | `int` | `0` | Data-parallel rank of this process |
| `ATOM_DP_RANK_LOCAL` | `int` | `0` | Local data-parallel rank (for SPMD mode) |
| `ATOM_DP_SIZE` | `int` | `1` | Total number of data-parallel groups |
| `ATOM_DP_MASTER_IP` | `str` | `"127.0.0.1"` | IP address of the data-parallel master |
| `ATOM_DP_MASTER_PORT` | `int` | `29500` | Port of the data-parallel master |
| `ATOM_DYNAMIC_CHUNKING_SMOOTH_FACTOR` | `float` | `0.75` | Default for `dynamic_chunking_smooth_factor`: `0` keeps the initial chunk, `1` uses the equal-latency chunk |
| ~~`ATOM_ENFORCE_EAGER`~~ | | | Removed. Use CLI flag `--enforce-eager` instead. |
| `ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION` | `bool` | `False` | Enable QK-norm + RoPE + cache + quant fusion for Qwen3 dense and MoE models |
| `ATOM_USE_TRITON_GEMM` | `bool` | `False` | Use Triton-based GEMM kernels instead of default backends |
| `ATOM_USE_TRITON_MXFP4_BMM` | `bool` | `False` | Use Triton-based MXFP4 batched matrix multiply |
| `ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION` | `bool` | `True` | Enable fused input RMSNorm + quantization for DeepSeek models |
| `ATOM_ENABLE_DS_QKNORM_QUANT_FUSION` | `bool` | `True` | Enable fused QK-norm + quantization for DeepSeek models |
| `ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION` | `bool` | `True` | Enable fused all-reduce + RMSNorm kernel |
| `ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_RMSNORM_QUANT` | `bool` | `True` | Enable AITER Triton fused RMSNorm + quantization for LLaMA models |
| `ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_SILU_MUL_QUANT` | `bool` | `True` | Enable AITER Triton fused SiLU + multiply + quantization for LLaMA models |

### Additional environment variables (used outside `envs.py`)

| Variable | Type | Default | Where Used | Description |
|---|---|---|---|---|
| `ATOM_TORCH_PROFILER_DIR` | `str` | `None` | `atom/config.py` (`Config.torch_profiler_dir`) | Directory for PyTorch profiler output; sets the default for `Config.torch_profiler_dir` |
| `ATOM_PROFILER_MORE` | `str` | `"0"` | `atom/model_engine/model_runner.py` | Set to `"1"` to enable detailed profiling (`record_shapes`, `with_stack`, `profile_memory`) |
| `HF_TOKEN` | `str` | `None` | `atom/config.py` (`get_hf_config`) | HuggingFace authentication token for gated model downloads |

## Decision tree — choosing a compilation level

```
Start
  |
  v
Is this a debugging / development run?
  |-- Yes --> Level 0 (NO_COMPILATION) or --enforce-eager
  |
  v
Do you need torch.compile but no graph splitting?
  |-- Yes, one-shot compile --> Level 2 (DYNAMO_ONCE)
  |-- Yes, keep Dynamo default --> Level 1 (DYNAMO_AS_IS)
  |
  v
Production inference on ROCm/HIP GPU?
  |-- Yes --> Level 3 (PIECEWISE) [default in EngineArgs]
              - Auto-sets CUDAGraphMode.TR_HYBRID
              - Auto-populates splitting_ops for attention ops
              - Pair with --cudagraph-capture-sizes for your batch profile
  |
  v
Need maximum decode throughput?
  |-- Yes --> set cudagraph_mode to FULL
              (a whole-forward request graph per decode bucket)
```

**Rules of thumb:**

- **Level 3** is the default for `EngineArgs` and is recommended for most
  production workloads.
- **Level 0** / `--enforce-eager` is useful for debugging, profiling, or when
  CUDA graphs are incompatible with your model.
- Match `--cudagraph-capture-sizes` to your expected batch sizes for optimal
  memory usage and launch latency.
- When using `--enable-dp-attention` or Expert Parallelism (`--enable-expert-parallel`),
  level 3 is still recommended.

## Source files

| File | Description |
|---|---|
| `atom/config.py` | `Config`, `CompilationConfig`, `CompilationLevel`, `CUDAGraphMode`, `QuantizationConfig`, `ParallelConfig`, `SpeculativeConfig`, `KVCacheTensor`, `KVCacheConfig`, `get_hf_config` |
| `atom/utils/envs.py` | All `ATOM_*` environment variable definitions with lazy evaluation |
| `atom/model_engine/arg_utils.py` | `EngineArgs` dataclass and CLI argument parser |
| `atom/sampling_params.py` | `SamplingParams` dataclass |
| `atom/model_engine/model_runner.py` | Uses `ATOM_PROFILER_MORE` and `ATOM_TORCH_PROFILER_DIR` for profiling |