ATOM model operations guide

ATOM (AiTer Optimized Model) wraps AITER kernels with model-level abstractions for LLM inference on AMD ROCm/HIP GPUs. This guide documents every operator class in atom/model_ops/, their AITER kernel mappings, quantization paths, and fused kernel chains.

Quick reference

ATOM Class	File	AITER Kernel / Import	Purpose
LinearBase	linear.py	tgemm.mm, gemm_a8w8, gemm_a8w8_bpreshuffle, gemm_a8w8_blockscale_bpreshuffle, gemm_a4w4	Quantized linear dispatch
ColumnParallelLinear	linear.py	(inherits LinearBase)	Column-sharded TP linear
RowParallelLinear	linear.py	(inherits LinearBase)	Row-sharded TP linear
QKVParallelLinear	linear.py	(inherits ColumnParallelLinear)	Fused Q/K/V projection
MergedColumnParallelLinear	linear.py	(inherits LinearBase)	Merged gate+up projection
Attention	base_attention.py	unified_attention_with_output_base (custom op)	Unified attention entry
MHA Attention	attention_mha.py	flash_attn_varlen_func, pa_fwd_asm, pa_persistent_fwd, pa_decode_gluon	Multi-head attention
MLAAttention	attention_mla.py	mla_decode_fwd, mla_prefill_fwd, concat_and_cache_mla, fused_qk_rope_concat_and_cache_mla	Multi-head latent attention
Shared DSA indexer scorer	indexer/	deepgemm_fp8_paged_mqa_logits, fp8_mqa_logits, flydsl_pa_mqa_logits_fp4, pa_mqa_logits_mxfp4, top_k_per_row_decode, top_k_per_row_prefill	Scoring and top-k for every DSA indexer (DeepSeek-V3.2 / GLM-5.x, GLM-5.3 kpool, DeepSeek-V4, V4.1): the attention builder publishes a step’s rows once (attn_metadata.index_rows: fp8_rows, fp4_rows, concat_rows, a DCP or TP rank’s share by RowSplit) cut to its logits plane (LogitsWorkspace), and a model calls score_topk against its cache (IndexCache)
DSA FP4 index cache	sparse_indexer_fp4.py	indexer_qk_rope_quant_and_cache	FP4 (E2M1 + e8m0) index-cache geometry and the DSA FP4 predicate
Fp4MqaRaggedMetadata	fp4_mqa_ragged_metadata.py	make_fp4_mqa_plan	FP4 MQA scoring’s ragged rows (DeepSeek-V4, V4.1, V3.2 decode and its DCP shard): rows grouped into sequences (query_start_loc, max_qlen, block tables, pages_per_block), each row’s bound, and their shape-only launch plan
FusedMoE	moe.py	aiter.fused_moe.fused_moe, asm_moe	Mixture of experts
RMSNorm	layernorm.py	rmsnorm2d_fwd, rmsnorm2d_fwd_with_add, fused_add_rmsnorm_pad	RMS normalization
LayerNorm	layernorm.py	layernorm2d_fwd, layernorm2d_fwd_with_add	Layer normalization
SiluAndMul	activation.py	aiter.silu_and_mul	SiLU gated activation
VocabParallelEmbedding	embed_head.py	F.embedding + TP all-reduce	Vocab embedding
ParallelLMHead	embed_head.py	tgemm.mm + tensor_model_parallel_all_gather	LM output head
RotaryEmbedding	rotary_embedding.py	aiter.rope_cached_positions_2c_fwd_inplace	Rotary position embedding
Sampler	sampler.py	aiter.mixed_sample_outer_exponential, aiter.ops.triton.topk.topk, aiter.ops.triton.softmax.softmax	Token sampling
RejectionSampler	rejection_sampler.py	Triton rejection_greedy_sample_kernel	Speculative decoding
constrain_logits, advance_states	required_channel.py	Triton _constrain_logits_kernel, _advance_states_kernel	tool_choice="required" channel mask (Kimi-K3)

AITER integration overview

ATOM is a thin model-level inference engine. Every compute-heavy operation delegates to an AITER kernel. The general pattern is:

1. An ATOM nn.Module owns model weights and configuration.

2. Its forward() method selects the appropriate AITER function based on quantization type, parallelism settings, and phase (prefill vs. decode).

3. Results are optionally reduced across tensor-parallel (TP) or data-parallel (DP) groups.

AITER kernel mapping table

ATOM Wrapper	AITER Function / Import Path	Backend Type
LinearBase.forward (No quant)	aiter.tuned_gemm.tgemm.mm	hipBLASLt
LinearBase.forward (per_Tensor FP8)	aiter.tuned_gemm.tgemm.mm with scales	hipBLASLt
LinearBase.forward (per_Token INT8)	aiter.gemm_a8w8	CK
LinearBase.forward (per_Token FP8)	aiter.gemm_a8w8_bpreshuffle	CK
LinearBase.forward (per_1x128 FP8)	aiter.gemm_a8w8_blockscale_bpreshuffle	CK
LinearBase.forward (per_1x32 MXFP4)	aiter.gemm_a4w4	CK
MHA prefill	aiter.flash_attn_varlen_func	ASM / CK
MHA decode (ASM)	aiter.pa_fwd_asm	ASM
MHA decode (persistent ASM)	aiter.pa_persistent_fwd	ASM
MHA decode (Triton)	aiter.ops.triton.gluon.pa_decode_gluon	Triton
MHA prefill (Triton unified)	aiter.ops.triton.unified_attention.unified_attention	Triton
MLA decode	aiter.mla.mla_decode_fwd	ASM
MLA prefill	aiter.mla.mla_prefill_fwd	ASM
MLA KV cache	aiter.concat_and_cache_mla	CK
RoPE	aiter.rope_cached_positions_2c_fwd_inplace	Triton
RMSNorm	aiter.rmsnorm2d_fwd	CK
SiLU+Mul	aiter.silu_and_mul	CK
TopK routing	aiter.topk_softmax, aiter.grouped_topk, aiter.biased_grouped_topk	CK
Sampling	aiter.mixed_sample_outer_exponential	CK
FusedMoE	aiter.fused_moe.fused_moe	CK
ASM MoE	aiter.fused_moe_bf16_asm.asm_moe	ASM
Quantization	aiter.get_hip_quant(QuantType)	CK / Triton

Linear operations

All linear layers inherit from LinearBase in atom/model_ops/linear.py.

Class hierarchy

LinearBase (nn.Module)
  +-- ReplicatedLinear          # No TP sharding
  |     +-- MergedReplicatedLinear
  +-- ColumnParallelLinear      # tp_dim=0, shard output
  |     +-- QKVParallelLinear   # Fused Q/K/V with per-head sharding
  +-- MergedColumnParallelLinear # tp_dim=0, merged gate+up
  +-- RowParallelLinear          # tp_dim=1, shard input, optional all-reduce

Quantization dispatch

LinearBase.forward() dispatches to different GEMM kernels based on QuantType:

QuantType	Weight dtype	GEMM Kernel	Scale Shape
No	BF16/FP16	tgemm.mm (hipBLASLt)	None
per_Tensor	FP8	tgemm.mm with scale_a, scale_b	[num_partitions, 1]
per_Token (INT8)	INT8	gemm_a8w8	[output_size, 1]
per_Token (FP8)	FP8	gemm_a8w8_bpreshuffle	[output_size, 1]
per_1x128	FP8	gemm_a8w8_blockscale_bpreshuffle	[ceil(N/128), ceil(K/128)]
per_1x32	MXFP4 (fp4x2)	gemm_a4w4 (preshuffled: K a multiple of 256, see MXFP4 K alignment)	[N, ceil(K/32)] (e8m0)

When x_scale is not provided, the input is dynamically quantized via get_hip_quant(quant_type).

Tensor parallel sharding

• ColumnParallelLinear (tp_dim=0): Shards weight rows (output dimension) across GPUs. Each GPU owns output_size / tp_size rows.

• RowParallelLinear (tp_dim=1): Shards weight columns (input dimension). If reduce_results=True, output is all-reduced across TP group.

• QKVParallelLinear: Extends ColumnParallelLinear with per-head sharding. Q heads are evenly divided; KV heads are either divided or replicated when num_kv_heads < tp_size.

• MergedColumnParallelLinear: Handles gate and up projections merged into a single weight with output_sizes as a list (e.g., [intermediate_size, intermediate_size]).

Weight processing

After loading, process_weights_after_loading() handles:

• e4m3fn to e4m3fnuz normalization (AMD FP8 format conversion).

• Weight reshuffling via shuffle_weights() for pre-shuffled GEMM kernels.

• Scale reshuffling via fp4_utils.e8m0_shuffle() for MXFP4 block scales.

• MXFP4 K alignment: a preshuffled MXFP4 weight must have K a multiple of 256 (MXFP4_PRESHUFFLE_K_ALIGN); the preshuffled GEMMs walk K in 256-wide steps and do not mask a tail. A gate_up -> act -> down MLP whose per-rank intermediate is unaligned pads it as a pair, like MoE’s intermediate_pad: pass mxfp4_k_pad_align(quant_config, down_prefix) as output_pad_align to MergedColumnParallelLinear and input_pad_align to RowParallelLinear. Only a pair that is MXFP4 or bf16 in the checkpoint is padded (a bf16 pair may be quantized online; its zeroed rows quantize to zero): the padded parameters are zeroed when allocated, and every load path writes the real shard at the front of its padded slot (requires act(0, 0) == 0). A padded layer requantized online from a quantized source (NVFP4, FP8), or a static pair whose halves aren’t both MXFP4 or bf16, is refused at construction. Any preshuffled MXFP4 weight still unaligned is refused here at load. Exempt: the gfx1250 32x4 scale layout, whose kernels take K % 128 (mxfp4_k_pad_align pads to 128 there), and ATOM_USE_FP4_NON_SHUFFLE_TRITON_GEMM=1, whose GEMM reads the raw layout. A refused layer can be kept out of MXFP4 (online quant: exclude_layer).

• Per-tensor requantization via requantize_with_max_scale() when multiple output partitions have separate scales.

Attention operations

Base: Attention (base_attention.py)

The top-level Attention class in base_attention.py is a dispatcher. It:

1. Selects the backend via get_attn_backend() from atom/utils/selector.py.

2. Instantiates the backend’s implementation class (impl_cls).

3. Registers itself in compilation_config.static_forward_context under layer_name.

4. On forward(), calls torch.ops.aiter.unified_attention_with_output_base, which is a custom op decorated with @mark_spliting_op — this prevents torch.compile from tracing into attention internals, enabling full-graph capture.

Backend selection logic (in selector.py):

Condition	Backend Class	Implementation
use_mla=True	AiterMLABackend	MLAAttention from attention_mla.py
use_mla=False	AiterBackend	Attention from attention_mha.py

Multi-head attention (attention_mha.py)

The MHA Attention class handles standard models (Llama, Qwen3, Mixtral, etc.).

Forward flow:

1. Reshape Q, K, V to [num_tokens, num_heads, head_dim].

2. Apply RoPE + KV cache write via rope_cache().

3. Dispatch to the appropriate backend via dispatch_backend().

RoPE + KV cache paths:

Condition	Kernel Chain
q_norm + k_norm + rotary_emb present	fused_qk_norm_rope_cache_quant_shuffle (single fused kernel for QK norm, RoPE, cache write, optional FP8 quant)
Triton path (sliding_window != -1 or head_dim != 128) + rotary_emb	fused_qk_rope_reshape_and_cache (Triton fused RoPE + reshape + cache)
ASM path + rotary_emb	rotary_emb(position, q, k) then reshape_and_cache or reshape_and_cache_with_pertoken_quant

Attention dispatch:

Phase	Condition	Method	AITER Kernel
Prefill	Always	prefill_attention	aiter.flash_attn_varlen_func
Decode	ATOM_USE_UNIFIED_ATTN and block_size == 256	paged_attention_persistent_asm	aiter.pa_persistent_fwd
Decode	past a paged kernel’s envelope, or ATOM_USE_UNIFIED_ATTN / use_flash_layout	paged_attention_unified	aiter.ops.triton.unified_attention
Decode	use_triton_attn=True	paged_attention_triton	torch.ops.aiter.pa_decode_gluon
Decode	Default	paged_attention_asm	aiter.pa_fwd_asm

The use_triton_attn flag is set when sliding_window != -1 or head_dim != 128.

Both paged kernels stop short of unified, at different points, and drafting multiplies the query group into both limits (base_attention.py):

Kernel	Envelope	Past it
gluon	query_length <= 4 and next_pow2(qlen) * max(16 // next_pow2(qlen), next_pow2(ratio)) <= 64, where ratio = q_heads / kv_heads	no layout arm; Triton fails to compile
pa_fwd_asm	qlen x q_heads/kv_heads <= 16	get_heuristic_kernel silently re-runs with mtp=1

_dispatch_decode is the single place that answers this; the runners do not re-decide. Unified carries one descale for the whole tensor, so a per-token quantized layer routed there raises rather than returning wrong numbers.

Multi-head latent attention (attention_mla.py)

MLAAttention implements DeepSeek’s MLA with a compressed KV representation. Key data structures:

@dataclass
class MLAModules:
    q_lora_rank: Optional[int]
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    qk_head_dim: int
    v_head_dim: int
    rotary_emb: torch.nn.Module
    q_proj: Optional[torch.nn.Module]
    kv_b_proj: torch.nn.Module
    o_proj: torch.nn.Module
    indexer: Optional[torch.nn.Module]

Forward flow:

1. If prefill and not sparse: Standard MHA-style prefill with flash_attn_varlen_func, preceded by kv_b_proj GEMM to produce K_nope and V from compressed kv_c_normed.

2. Otherwise: Fused Q projection + K up-projection via batched FP8/FP4 BMM (_q_proj_and_k_up_proj), then:

   • fused_qk_rope_concat_and_cache_mla writes to KV cache.

   • Decode: mla_decode_fwd (ASM persistent MLA kernel).

   • Prefill (sparse): mla_prefill_fwd.

3. V up-projection + O projection via batched BMM (_v_up_proj_and_o_proj).

Batched GEMM backends for MLA projections:

Condition	Kernel
ATOM_USE_TRITON_MXFP4_BMM=True	batched_gemm_a16wfp4 (Triton FP4 BMM)
Default	batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant (Triton FP8 BMM)

Prefill GEMM optimizations (for kv_b_proj):

Condition	Kernel
ATOM_USE_TRITON_GEMM=True + FP4 weights	fused_gemm_afp4wfp4_preshuffle_split_cat (GEMM + split K/V + cat rope in one kernel)
ATOM_USE_TRITON_GEMM=True + FP8 weights	fused_gemm_a8w8_blockscale_preshuffle_split_cat
Default	kv_b_proj(kv_c_normed) then manual split + cat

Backend abstraction (attentions/backends.py)

The AttentionBackend abstract class defines three required methods:

• get_name() — Returns backend identifier string.

• get_builder_cls() — Returns the AttentionMetadataBuilder subclass.

• get_impl_cls() — Returns the attention implementation class.

CommonAttentionBuilder provides shared metadata preparation (slot mapping, block tables, cumulative sequence lengths) used by both AiterBackend and AiterMLABackend. It marshals block_tables on every prefill step, whether or not that step also uploads the buffer, because the slot arithmetic in token_layout/prefill.py reads the packed table back rather than the batch’s ragged rows.

KV cache operations

Operation	AITER Kernel	Used By
Standard KV cache write	aiter.reshape_and_cache	MHA (BF16 KV)
FP8 KV cache write	aiter.reshape_and_cache_with_pertoken_quant	MHA (FP8 KV)
MLA KV cache write	aiter.concat_and_cache_mla	MLA prefill
Fused QK RoPE + MLA cache	aiter.fused_qk_rope_concat_and_cache_mla	MLA decode

Mixture of experts (MoE)

FusedMoE class (moe.py)

FusedMoE is the top-level MoE module. It handles:

• Expert routing via select_experts().

• Weight creation and quantization dispatch via quant_method.

• Tensor/Expert/Data parallelism via FusedMoEParallelConfig.

• Optional shared expert fusion and MORI communication.

Constructor parameters:

FusedMoE(
    num_experts: int,        # Global number of experts
    top_k: int,              # Experts per token
    hidden_size: int,        # Input hidden dimension
    intermediate_size: int,  # Expert intermediate dimension
    reduce_results: bool,    # Whether to all-reduce output
    renormalize: bool,       # Renormalize routing weights
    use_grouped_topk: bool,  # Use grouped top-k (DeepSeek)
    activation: ActivationType,  # Silu, Gelu, Swiglu, etc.
    ...
)

Quantization methods

FusedMoE selects a quant_method at construction time:

Quant Config	Method Class	GEMM Kernel
QuantType.No	UnquantizedFusedMoEMethod	aiter.fused_moe.fused_moe
FP8 (dtypes.fp8)	Fp8MoEMethod	aiter.fused_moe.fused_moe with quant_type
FP8 compressed-tensors	CompressedTensorsFp8MoEMethod	aiter.fused_moe.fused_moe or asm_moe
MXFP4 (dtypes.fp4x2)	Mxfp4MoEMethod	aiter.fused_moe.fused_moe or Triton triton_kernel_fused_experts

The ASM MoE path (asm_moe from aiter.fused_moe_bf16_asm) is used by FP8 methods and supports a16 mode where activations remain in BF16/FP16 while weights are FP8/INT8.

TopK routing (topK.py)

Routing Function	AITER Kernel	Used For
rocm_aiter_topk_softmax	aiter.topk_softmax	Standard top-k (Mixtral)
rocm_aiter_grouped_topk	aiter.grouped_topk	Grouped top-k (DeepSeek)
rocm_aiter_biased_grouped_topk	aiter.biased_grouped_topk	Biased grouped top-k (DeepSeek V3)

Shared expert fusion: When is_rocm_aiter_fusion_shared_expert_enabled() returns True, the top-k buffers are extended with shared expert IDs appended after routed expert IDs. This allows shared expert computation to be fused into the same MoE kernel call. The metadata is initialized via init_aiter_topK_meta_data(). Under DP attention with MoRI EP, fusion is wired up only for MXFP4 / NVFP4 routed experts; FP8 MoE models such as DeepSeek-V3 must run with ATOM_FUSE_SHARED_EXPERT=0, or FusedMoE raises NotImplementedError.

FusedMoEParallelConfig

@dataclass
class FusedMoEParallelConfig:
    tp_size: int       # Tensor parallel size
    dp_size: int       # Data parallel size
    ep_size: int       # Expert parallel size
    tp_rank: int
    dp_rank: int
    ep_rank: int
    use_ep: bool       # Whether expert parallelism is active
    local_ep_size: int # Local EP size (GPUs per node * TP)

Key properties:

• use_all2all_kernels: True when dp_size > 1, EP is enabled, and MORI is available.

• use_mori_kernels: Always True (currently).

MORI integration (fused_moe/mori_prepare_finalize.py)

MORI (MoE Router Infrastructure) provides all-to-all communication kernels for expert parallelism. MoriPrepareAndFinalize implements:

• prepare(): Dispatches tokens to remote experts via mori_op.dispatch(), in the layer’s routed-rows wire (MXFP8 with e8m0 scales where the experts take MXFP8, ATOM_MORI_FP4_DISPATCH for MXFP4, bf16 otherwise).

• finalize(): Combines expert outputs via mori_op.combine() and copies results back.

The FusedMoEModularKernel orchestrates the prepare-compute-finalize pipeline.

Routing before the rows move (fused_moe/routed_rows.py)

A MoE method splits into select_experts(layer, x, router_logits) (logical routes), place_experts(layer, topk_weights, topk_ids, gathered) (the EPLB remap and load record, the shared column in the layer’s expert layout) and apply(layer, x, topk_weights, topk_ids, ..., x_scale=None). FusedMoE routes on the rank that owns the rows, before any of them move: the MORI dispatch, the DP gather (DP attention without the all2all kernels), the PCP gather (ATOM_PCP_MOE_MERGE, inside FusedMoE._run_experts) and the comm-fused DPA gather all carry top-k weights and ids rather than router logits. Rows are placed after every gather, so a gathered row – held alike by every rank – picks the same EPLB replica on each (gathered); and where the legacy layout widens in the prefilled shared column (LEGACY_AITER), whose ids say by row position which rank runs a row’s shared expert, a gather is followed by reassign_gathered_shared_ids, which gives the rows this rank’s ids for their gathered positions. Every routing kernel takes the pass’s real row count, Context.scheduled_tokens_gpu (an entry of the step’s forward_vars["batch_info"], read on device so a replayed graph follows its own step), and routes the rows past it to -1, so they are neither sent nor computed. Only a consumer of the pass’s whole rows (running_tokens) gets the count (pass_scheduled_tokens_gpu): the MoE sort reads it as its row count, so a subset of the rows – a PCP shard, gathered rows – gets None and runs every row, unless it declares a pass of its own with its own count (ForwardContext.use_pass_rows, as V4.1’s decoder does over a prefill’s tails with ReplayPlan.filled_gpu / finishing_gpu). The rows travel in RoutedRowsWire, the method’s routed_rows_wire – the one format both the gathers and the MoRI v1 dispatch use: MXFP4 under ATOM_MORI_FP4_DISPATCH, MXFP8 where the experts take MXFP8 activations (FlyDSL fused_moe, the Triton a8w4 EP experts), bf16 otherwise (the TP Triton path, MegaMoE, the MoRI v2 transport). Rows already in the wire carry their e8m0 rows as x_scale through apply, the modular kernel and prepare, so nothing is quantized twice.

MoE quantization config (fused_moe/config.py)

FusedMoEQuantConfig describes activation and weight quantization for MoE layers:

@dataclass
class FusedMoEQuantConfig:
    _a1: FusedMoEQuantDesc   # First activation (input to gate_up)
    _a2: FusedMoEQuantDesc   # Second activation (input to down_proj)
    _w1: FusedMoEQuantDesc   # gate_up_proj weights
    _w2: FusedMoEQuantDesc   # down_proj weights

Factory functions:

• fp8_w8a8_moe_quant_config() — FP8 weights and activations.

• mxfp4_w4a16_moe_quant_config() — MXFP4 weights, unquantized activations.

• FUSED_MOE_UNQUANTIZED_CONFIG — No quantization.

Triton MoE fallback (fused_moe_triton.py)

triton_kernel_fused_experts() provides a Triton-based MoE path. It takes rows routed already (FusedMoE.select_experts, before any gather or dispatch): routing_from_dispatched() sorts their top-k ids into the Triton routing data, with a gate_valid mask for rows another rank owns or rows past the real row count, and the expert GEMMs run on that. This path is used for MXFP4 MoE under ATOM_USE_TRITON_MOE.

Normalization

RMSNorm (layernorm.py)

RMSNorm supports multiple forward paths depending on configuration flags:

Condition	Kernel / Path	Returns
x_pad_to_multiple > 0, no residual	fused_rmsnorm_pad_ (Triton fused_add_rmsnorm_pad)	Padded output
x_pad_to_multiple > 0, with residual	fused_add_rmsnorm_pad_	(output, residual)
fused_allreduce=True and tp_size > 1	tensor_model_parallel_fused_allreduce_rmsnorm	(output, residual)
fused_quant=True and x_scale provided	fused_rms_fp8_per_tensor_static_quant	(FP8 output, scale)
fused_quant=True and per_1x32	fused_rms_mxfp4_quant	(MXFP4 output, scale)
Default, no residual	rmsnorm2d_fwd	Output
Default, with residual	rmsnorm2d_fwd_with_add	(output, residual)

Constructor parameters:

RMSNorm(
    dim: int,
    eps: float = 1e-6,
    x_pad_to_multiple: int = 0,
    fused_allreduce: bool = False,
    fused_quant: bool = False,
    quant_config: Optional[QuantizationConfig] = None,
)

LayerNorm (layernorm.py)

LayerNorm wraps layernorm2d_fwd and layernorm2d_fwd_with_add (with bias support):

LayerNorm(dim: int, eps: float = 1e-6)

• Without residual: layernorm2d_fwd(x, weight, bias, eps)

• With residual: layernorm2d_fwd_with_add(out, x, residual, residual_out, weight, bias, eps)

Activation functions

SiluAndMul (activation.py)

SiluAndMul computes SiLU(x_first_half) * x_second_half. It splits the last dimension in half.

Condition	Kernel	Output
fused_quant=True + x_scale provided (FP8)	fused_silu_mul_fp8_per_tensor_static_quant	(FP8 output, scale)
fused_quant=True + per_1x32 (MXFP4)	fused_reduce_act_mul_and_mxfp4_quant (via mxfp4_act_mul_quant_fuse)	(MXFP4 output, scale)
Default	aiter.silu_and_mul(out, x)	BF16 output

Constructor:

SiluAndMul(
    fused_quant: bool = False,
    quant_config: Optional[QuantizationConfig] = None,
)

Embedding and output head

VocabParallelEmbedding (embed_head.py)

Partitions the vocabulary across TP ranks. Each rank holds num_embeddings / tp_size rows.

Forward:

1. Mask input token IDs to this rank’s partition range [vocab_start_idx, vocab_end_idx).

2. F.embedding() on local partition.

3. Zero out out-of-range positions.

4. all_reduce() across TP group.

ParallelLMHead (embed_head.py)

Extends VocabParallelEmbedding for the output projection. Key differences:

• Forward extracts only the last token per sequence during prefill (via cu_seqlens_q[1:] - 1).

• Uses tgemm.mm(x, self.weight, self.bias) for the logit computation (not F.linear).

• Calls tensor_model_parallel_all_gather() to gather logits across TP ranks.

Rotary position embedding (RoPE)

RotaryEmbedding (rotary_embedding.py)

Precomputes cos/sin caches at initialization and applies RoPE in-place.

Constructor:

RotaryEmbedding(
    head_size: int,
    rotary_dim: int,
    max_position_embeddings: int,
    base: float,
    is_neox_style: bool = True,
    dtype: Optional[torch.dtype] = None,
)

Forward: Calls aiter.rope_cached_positions_2c_fwd_inplace(query_, key_, cos, sin, positions, rotate_style, ...) which applies RoPE to Q and K tensors in-place using precomputed caches indexed by position IDs.

get_rope() factory

get_rope(head_size, rotary_dim, max_position, base, rope_scaling=None)

Returns a cached RotaryEmbedding instance. Currently rope_scaling must be None.

Integration in attention

• MHA (attention_mha.py): RoPE is applied during the rope_cache() phase, either via the fused fused_qk_norm_rope_cache_quant_shuffle kernel, via fused_qk_rope_reshape_and_cache, or via standalone rotary_emb(position, q, k).

• MLA (attention_mla.py): RoPE is applied to q_pe and k_rope tensors. During decode, this is fused into fused_qk_rope_concat_and_cache_mla. During prefill, it is applied via self.rotary_emb(positions, prefill_q_pe, k_rope).

Sampling

Sampler (sampler.py)

Unified sampling supporting both greedy (temperature=0) and random (temperature>0) sampling in a single kernel call.

Forward:

def forward(self, logits, temperatures) -> sampled_tokens:
    mixed_sample_outer_exponential(sampled_tokens, logits, exponential, temperatures, eps)

aiter.mixed_sample_outer_exponential performs temperature-scaled exponential sampling: it divides logits by temperature, then uses the Gumbel-max trick with pre-generated exponential random variates.

Fallback methods (currently unreachable due to early return):

• greedy_sample(): aiter.ops.triton.topk.topk(logits, 1)

• random_sample(): aiter.ops.triton.softmax.softmax(logits) followed by exponential sampling and topk.

Required channel (required_channel.py)

Holds a request with a RequiredChannel (atom/sampling_params.py; today only Kimi-K3 under tool_choice="required") to its tools channel by masking the target logits before any sampler reads them. Until the channel opens, the end ids (the format’s own plus the engine’s EOS and stop_token_ids) are -inf; once the output ends with after (K3: <|close|>response<|sep|>), each next row keeps only the forced id of opening. The channel is open once opening is written, or once the model writes opening[:opens_at] itself (K3: <|open|>tools<|sep|>); the state is then UNCONSTRAINED and the request costs nothing more.

• constrain_logits masks a step’s rows in place: plain decode and prefill rows by the request’s state, verification rows as if the drafts before them were accepted, and the bonus row after every draft. A mismatching draft is therefore rejected by the ordinary rejection sampler.

• advance_states folds the step’s emitted ids into each state on the GPU, so a deferred-output step reads its predecessor’s states through decode_src before the host has seen those ids.

constrain_logits runs VOCAB_PROGRAMS (16) programs per request, each writing every 16th 1,024-column block of a forced row without reading it; advance_states runs one per request. On MI355X over K3’s vocabulary a step with no forced row costs what one program per request did (5-7 us for a DSpark7 verify step at bs 1-64), and one with all seven draft rows forced 15-52 us. A batch with no unopened channel launches neither kernel. constrain_logits_reference / advance_states_reference are the pure-torch twins the tests hold the kernels to.

RejectionSampler (rejection_sampler.py)

Implements rejection sampling for speculative decoding (MTP). Given draft token IDs and target model logits:

1. Computes target_argmax = target_logits.argmax(dim=-1).

2. Runs a Triton kernel rejection_greedy_sample_kernel that sequentially compares draft tokens against target argmax, accepting until first mismatch.

3. On full acceptance, appends the bonus token.

4. Returns (output_token_ids, num_bonus_tokens).

Fused kernel chains

ATOM uses fused kernels to reduce memory traffic by combining multiple operations into a single kernel launch.

Fused Operation	Components	Controlled By	AITER Kernel
RMSNorm + FP8 quant	RMSNorm, per-tensor FP8 static quant	RMSNorm(fused_quant=True) + x_scale	fused_rms_fp8_per_tensor_static_quant
RMSNorm + MXFP4 quant	RMSNorm, per-1x32 MXFP4 quant	RMSNorm(fused_quant=True) + QuantType.per_1x32	fused_rms_mxfp4_quant
RMSNorm + add + pad	Residual add, RMSNorm, output padding	RMSNorm(x_pad_to_multiple>0)	fused_add_rmsnorm_pad
AllReduce + RMSNorm	TP all-reduce, RMSNorm	RMSNorm(fused_allreduce=True)	tensor_model_parallel_fused_allreduce_rmsnorm
SiLU + mul + FP8 quant	SiLU activation, multiply, FP8 quant	SiluAndMul(fused_quant=True) + x_scale	fused_silu_mul_fp8_per_tensor_static_quant
SiLU + mul + MXFP4 quant	SiLU activation, multiply, MXFP4 quant	SiluAndMul(fused_quant=True) + QuantType.per_1x32	fused_reduce_act_mul_and_mxfp4_quant
QK norm + RoPE + cache + quant	Q/K norm, RoPE, KV cache write, optional FP8 quant, weight shuffle	q_norm + k_norm + rotary_emb all present	fused_qk_norm_rope_cache_quant_shuffle
RoPE + reshape + cache	RoPE, K reshape, KV cache write	Triton attention path	fused_qk_rope_reshape_and_cache
QK RoPE + MLA cache	Q RoPE, KV concat, MLA cache write, FP8 quant	MLA decode path	fused_qk_rope_concat_and_cache_mla
GEMM + split + cat (FP4)	KV_b_proj GEMM, split K_nope/V, cat K_rope	ATOM_USE_TRITON_GEMM=True + FP4 weights	fused_gemm_afp4wfp4_preshuffle_split_cat
GEMM + split + cat (FP8)	KV_b_proj GEMM, split K_nope/V, cat K_rope	ATOM_USE_TRITON_GEMM=True + FP8 weights	fused_gemm_a8w8_blockscale_preshuffle_split_cat
FP8 BMM + RoPE + cache (MLA)	Batched FP8 BMM, RoPE, MLA KV cache write	MLA decode with FP8	fused_fp8_bmm_rope_cat_and_cache_mla
FP4 BMM + RoPE + cache (MLA)	Batched FP4 BMM, RoPE, MLA KV cache write	MLA decode with MXFP4	fused_fp4_bmm_rope_cat_and_cache_mla

Source files

atom/model_ops/

File	Description
linear.py	LinearBase, ColumnParallelLinear, RowParallelLinear, QKVParallelLinear, MergedColumnParallelLinear, ReplicatedLinear, MergedReplicatedLinear
activation.py	SiluAndMul with fused FP8/MXFP4 quantization
layernorm.py	RMSNorm, LayerNorm with fused allreduce/quant/pad variants
base_attention.py	Top-level Attention dispatcher with custom op registration
attention_mha.py	MHA implementation: prefill (flash), decode (ASM/Triton paged attention)
attention_mla.py	MLAAttention, MLAModules — DeepSeek MLA with compressed KV
moe.py	FusedMoE, FusedMoEParallelConfig, UnquantizedFusedMoEMethod, Fp8MoEMethod, Mxfp4MoEMethod, CompressedTensorsFp8MoEMethod
fused_moe_triton.py	triton_kernel_fused_experts — Triton MoE experts on rows routed already (routing_from_dispatched)
embed_head.py	VocabParallelEmbedding, ParallelLMHead
rotary_embedding.py	RotaryEmbedding, get_rope
topK.py	rocm_aiter_topk_softmax, rocm_aiter_grouped_topk, init_aiter_topK_meta_data
sampler.py	Sampler — unified greedy/random sampling
rejection_sampler.py	RejectionSampler — speculative decoding rejection sampling
required_channel.py	constrain_logits, advance_states — tool_choice="required" channel mask (Kimi-K3)
base_config.py	QuantizeMethodBase abstract class
utils.py	Helper utilities: shuffle_weights, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize, etc.

atom/model_ops/attentions/

File	Description
backends.py	AttentionBackend, AttentionMetadataBuilder, CommonAttentionBuilder, AttentionImpl abstract classes
aiter_attention.py	AiterBackend, AiterAttentionMetadataBuilder — MHA backend with persistent ASM paged attention support
aiter_mla.py	AiterMLABackend, AiterMLAMetadataBuilder — MLA backend with sparse attention support

atom/model_ops/attentions/pool_layout/ and .../token_layout/

Two packages, two axes. pool_layout answers where a byte lives in the cache pools and is a function of the config; token_layout answers where this step’s tokens go and is a function of the batch, so it is rebuilt every step. A backend consumes both and belongs to neither.

Membership has three parts, and only the third is machine-checkable. By topic — the axis decides, not the dependencies. Arithmetic, not staging — a member computes an answer and does not write or upload a forward_vars mirror; page_unit_geometry is a mixin over self.model_runner and still qualifies because it only reads, so taking self is not the line, having a side effect on the step’s buffers is. Reachable without aiter or the rest of atom, so a member is importable on a runner with no AITER build and no GPU — CI is such a runner, and one import failure during collection aborts the whole run rather than one test, so tests/test_layout_packages.py enforces that part over both packages. Do not mistake the checkable part for the rule.

Two absences are deliberate. v4_kernels/pool_index.py is v4_pool_geometry’s device-side half — the same row formulas as @triton.jit device functions for the eight kernels that address the pool — and stays with the kernels; tests/test_pool_index.py pins the two sides together. And a per-token shape that only one caller has stays with that caller: V4’s decode positions are ragged, and the one-token decode slot mapping comes from last_block_num_tokens rather than from a position.

File	Description
pool_layout/sub_pool_spec.py	SubPoolSpec, page_pool, state_pool, plan_pools — sub-pool sizing as arithmetic over a byte budget
pool_layout/v4_pool_geometry.py	UnifiedPoolGeometry, WindowParams, the compress ratios — where a DeepSeek-V4 row lives, and which rows a step may see
pool_layout/page_unit_geometry.py	PageUnitGeometryMixin — where a K3 checkpoint image’s bytes land in the MLA paged pool
pool_layout/entry_arena.py	EntryField, entry_bytes_for, plan_regions, EntryMajorArena — what one entry of a cache class holds, and where those bytes are put
pool_layout/paged_state_copy.py	plan_segmented_copy, launch_copy_descriptor — scattering that byte run across PAGE units and back
token_layout/prefill.py	prefill_positions — where a ragged prefill chunk’s tokens sit in their own sequences
token_layout/decode.py	decode_positions — the same for the rectangular speculative decode step
token_layout/slots.py	slot_mapping — which KV slot a token is written to, one gather for both sides
token_layout/batch_ids.py	build_batch_ids — builds the token → sequence map every kernel resolves per-sequence data through. Query lengths in gives batch_id_per_q_token, context lengths in gives batch_id_per_k_token; the two axes are not interchangeable

atom/model_ops/fused_moe/

File	Description
config.py	FusedMoEConfig, FusedMoEQuantConfig, FusedMoEQuantDesc, GroupShape, factory functions (fp8_w8a8_moe_quant_config, mxfp4_w4a16_moe_quant_config)
modular_kernel.py	FusedMoEModularKernel, FusedMoEPrepareAndFinalize, ExpertTokensMetadata — modular MoE kernel pipeline
mori_prepare_finalize.py	MoriPrepareAndFinalize — MORI all-to-all dispatch/combine for expert parallelism
routed_rows.py	RoutedRowsWire, resolve_routed_rows_wire — the format routed rows travel in across a dispatch or a DP / PCP gather
utils.py	MoE utility functions

atom/utils/

File	Description
selector.py	get_attn_backend() — selects AiterBackend or AiterMLABackend based on use_mla flag