Struct ForwardContext 

pub struct ForwardContext<'a> {

    pub buffers: &'a BufferArena,
    pub hc_row_offset: usize,
    pub gpu: &'a dyn GpuBackend,
    pub config: &'a ModelConfig,
    pub dispatch: &'a GemmDispatch,
    pub derived: &'a DerivedWeights,
    pub levers: &'a ModelLevers,
    pub stats: &'a ModelStats,
    pub attn_metadata: Option<AttnMetadataDev>,
    pub profile: bool,
    pub comm: Option<&'a dyn CommBackend>,
    pub graph_capture: bool,
    pub decode_step: bool,
    pub gdn_exact_replay: bool,
    pub token_ids: Option<DevicePtr>,
    pub host_token_ids: Option<&'a [u32]>,
    pub routed_lora_layers: Option<&'a [Option<LoraLayerWeights>]>,
    pub midchunk_capture: Option<MidchunkCapture<'a>>,
    pub moe_lora_route: MoeLoraRoute,
}

Shared context for a single forward pass step.

Provides access to GPU, buffers, and config without coupling layer implementations to the model struct.

Fields

buffers: &'a BufferArena

Pre-allocated scratch buffers.

hc_row_offset: usize

mHC highway ROW offset for this pass (#753 item B, mixed steps): the fused decode+prefill step gives the prefill chunk highway rows at padded_n so they live disjoint from the decode rows, mirroring the hidden/residual layout. 0 everywhere else.

gpu: &'a dyn GpuBackend

GPU backend for kernel launches and memory ops.

config: &'a ModelConfig

Model configuration (dimensions, hyperparameters).

dispatch: &'a GemmDispatch

Which GEMM implementation each projection takes. Carried rather than read from a static so it cannot outlive the model whose flags it encodes — see layers::ops::GemmDispatch.

derived: &'a DerivedWeights

Re-encoded copies of this model’s weights, memoized for this model’s lifetime. Carried rather than kept in a static keyed by device pointer, where a recycled address would HIT after a model swap.

levers: &'a ModelLevers

Kernel-path levers for this model — the SSM/GDN variant, FFN routing, MoE quantization, LoRA mode, diagnostics. The non-GEMM half of the lever set; dispatch is the GEMM half.

stats: &'a ModelStats

This model’s diagnostic counters and one-shot dump latches. Carried for the same reason as levers: a counter that spans a model swap averages two models and describes neither, and a one-shot latch that already fired swallows the next model’s dump.

attn_metadata: Option<AttnMetadataDev>

Pre-uploaded attention metadata (None if no attention layers).

profile: bool

Profile mode: sync+time per-operation within layers.

comm: Option<&'a dyn CommBackend>

Communication backend for expert parallelism (EP) all-reduce. None when running single-GPU (no distributed communication).

graph_capture: bool

True when inside CUDA graph capture (between begin_capture/end_capture). MoE layers use sync all_reduce (capturable) instead of async (event-based).

decode_step: bool

True ONLY on the single-token decode step, where attn_metadata’s positions, slot, seq_len and block_table are the step’s SCALARS at stable addresses.

🪤 prefill_default drives a layer that has no prefill of its own by calling its decode once per token — with the PREFILL context, whose positions/slot are per-token ARRAYS and whose block_table/seq_len are NULL unless the pass is paged. A layer that reads those pointers as decode scalars gets an illegal address on the first prompt. Check this flag, not attn_metadata.is_some().

gdn_exact_replay: bool

True when this prefill pass continues from a restored Marconi SSM snapshot (warm prefix-cache hit). GDN layers must then take the bit-faithful WY4 recurrence instead of the FLA chunked kernel: FLA’s chunk grid is anchored at the (arbitrary) snapshot offset and its bf16 intermediates drift vs the pass that originally produced the cached K/V, and the replay range [snap_tok, matched) is rewritten into SHARED prefix-cache blocks — non-exact recompute poisons them and the drift ratchets across turns (2026-06-10 warm-hit stutter).

token_ids: Option<DevicePtr>

Device [num_tokens] u32 token IDs for the tokens being processed this pass, in the SAME order the per-token MoE loop visits them. Required by DeepSeek-V4 hash-MoE layers (static tid2eid[token_id] routing); None for models without hash routing. Must be a STABLE address across the layer loop (and, under CUDA-graph decode, uploaded before each replay).

host_token_ids: Option<&'a [u32]>

HOST copy of the same token ids, when the caller had them in hand (decode always does — it uploads token_ids FROM this value; chunked prefill likewise). PLE computes its n-gram ids on the host, and reading them back off the device costs a synchronous D2H per decode step — pure overhead, and capture-unsupported inside a CUDA graph.

routed_lora_layers: Option<&'a [Option<LoraLayerWeights>]>

#30 (routed-prefill precision): the REQUEST slot’s per-layer LoRA pairs, GLOBAL-layer-indexed (len == num_hidden_layers), set ONLY at the prefill entries and ONLY when the request routes to a NON-active slot. Some makes the K/V/O prefill apply sites select the request slot’s pair and fold it through the SAME dense apply_lora_delta (dense_gemm_tc) the ACTIVE adapter uses — numerically identical to serving that adapter active, instead of the per-row bgmv (whose fp accumulation order tips razor-margin tokens). None (active/base request, no LoRA, and every decode/verify/mtp/moe pass) leaves the installed-active-pair path byte-identical. Prefill runs eager (graph_capture: false) so this per-pass CPU borrow is safe.

midchunk_capture: Option<MidchunkCapture<'a>>

Default-ON mid-chunk SSM tail capture (opt-out ATLAS_SSM_TAIL_MIDCHUNK=0).

Some only on the single prefill pass whose local token range spans the block-floored matched-prefix boundary tb. GDN/SSM layers then split their recurrent (h_state) and conv (conv_state) kernels at cap_local and copy the @tb state into the reserved snapshot slot. None (default) => no split, byte-identical to prior behavior.

moe_lora_route: MoeLoraRoute

Feature-1 MoE-LoRA per-request fold decision for this forward pass, resolved by TransformerModel::moe_lora_route from the owning request’s adapter_slot. Governs the prefill router/expert fold hooks (layers/moe/lora.rs). Ignored when no MoE adapter is installed (self.lora == None short-circuits first — byte-identical off). Default Fold keeps legacy single-request call sites unchanged.