pub struct TpE4m3HostBounce { /* private fields */ }Expand description
Multi-context TP correctness runtime. Each rank owns an independent Engine and CUDA context.
Host bounce is the default oracle. Native P2P is opt-in and preserves the oracle’s global checkpoint-block reduction order; it remains a correctness path until serving gates and repeated performance evidence qualify it.
Implementations§
Source§impl TpE4m3HostBounce
impl TpE4m3HostBounce
pub fn new(devices: &[usize]) -> Result<Self, Box<dyn Error>>
pub fn new_native_p2p(devices: &[usize]) -> Result<Self, Box<dyn Error>>
pub fn new_native_p2p_device_arithmetic( devices: &[usize], ) -> Result<Self, Box<dyn Error>>
Sourcepub fn new_single_rank_oracle(device: usize) -> Result<Self, Box<dyn Error>>
pub fn new_single_rank_oracle(device: usize) -> Result<Self, Box<dyn Error>>
Single-rank execution of the canonical checkpoint-block TP program.
This is an oracle for distributed exactness, not a serving topology. It lets gates compare TP=1 and TP>1 with the same packing, kernel launches, and deterministic reduction order.
pub fn devices(&self) -> &[usize]
pub fn native_p2p(&self) -> bool
pub fn bulk_p2p(&self) -> bool
pub fn expert_activation_label(&self) -> &'static str
pub fn expert_accumulation_label(&self) -> &'static str
pub fn expert_output_label(&self) -> &'static str
pub fn transport_label(&self) -> &'static str
pub fn device_names(&self) -> Result<Vec<String>, Box<dyn Error>>
Sourcepub fn rank_engine(&self, rank: usize) -> Option<&Engine>
pub fn rank_engine(&self, rank: usize) -> Option<&Engine>
Correctness-gate access to the engine that owns one TP rank.
Model execution should prefer collective methods on this runtime. This accessor exists so focused gates can prove that the rank-local projection outputs remain device-resident through the next ownership boundary before that boundary is wired into serving.
pub fn allocate_tp_kv_cache( &self, kv_dim_k: usize, kv_dim_v: usize, capacity: usize, ) -> Result<ResidentTpKvCache, Box<dyn Error>>
pub fn allocate_tp_swa_kv_cache( &self, kv_dim_k: usize, kv_dim_v: usize, capacity: usize, window: usize, ) -> Result<ResidentTpKvCache, Box<dyn Error>>
pub fn grow_tp_kv_cache( &self, source: &ResidentTpKvCache, target_capacity: usize, rows: usize, ) -> Result<ResidentTpKvCache, Box<dyn Error>>
pub fn hydrate_tp_kv_cache( &self, cache: &mut ResidentTpKvCache, rows: usize, k_rows: &[u8], v_rows: &[u8], ) -> Result<(), Box<dyn Error>>
pub fn hydrate_tp_kv_cache_from( &self, cache: &mut ResidentTpKvCache, logical_len: usize, resident_start: usize, k_rows: &[u8], v_rows: &[u8], ) -> Result<(), Box<dyn Error>>
pub fn append_tp_kv_transaction( &self, cache: &mut ResidentTpKvCache, transaction: TpKvTransaction, k_shards: &[CudaSlice<f32>], v_shards: &[CudaSlice<f32>], rows: usize, ) -> Result<(), Box<dyn Error>>
Sourcepub fn append_tp_kv_transaction_inner(
&self,
cache: &mut ResidentTpKvCache,
transaction: TpKvTransaction,
k_shards: &[CudaSlice<f32>],
v_shards: &[CudaSlice<f32>],
rows: usize,
external_rank_appends: bool,
) -> Result<(), Box<dyn Error>>
pub fn append_tp_kv_transaction_inner( &self, cache: &mut ResidentTpKvCache, transaction: TpKvTransaction, k_shards: &[CudaSlice<f32>], v_shards: &[CudaSlice<f32>], rows: usize, external_rank_appends: bool, ) -> Result<(), Box<dyn Error>>
external_rank_appends: the dcw path already wrote the rank rows (device-counter
append) — run everything EXCEPT the per-rank quantize/append loop (plan validation,
rebase arm — unreachable when the caller peeked — and the absolute len-mirror sets,
which land the same value the in-stream inc produced).
pub fn commit_tp_kv_transaction( &self, cache: &mut ResidentTpKvCache, transaction: TpKvTransaction, accepted_rows: usize, ) -> Result<(), Box<dyn Error>>
Sourcepub fn commit_tp_kv_transaction_external(
&self,
cache: &mut ResidentTpKvCache,
transaction: TpKvTransaction,
accepted_rows: usize,
) -> Result<(), Box<dyn Error>>
pub fn commit_tp_kv_transaction_external( &self, cache: &mut ResidentTpKvCache, transaction: TpKvTransaction, accepted_rows: usize, ) -> Result<(), Box<dyn Error>>
Commit for the external-appends (token graph) path: host bookkeeping only, NO absolute len-mirror sets. The graph’s in-stream inc_i32 owns the device counters; a rank-stream set here has no ordering edge against the NEXT token’s graph launch (graph children do not wait on the rank streams), so it can land AFTER that graph’s inc and drag the counter backward mid-token.
pub fn rollback_tp_kv_transaction( &self, cache: &mut ResidentTpKvCache, transaction: TpKvTransaction, ) -> Result<(), Box<dyn Error>>
pub fn tp_kv_device_lengths( &self, cache: &ResidentTpKvCache, ) -> Result<Vec<i32>, Box<dyn Error>>
pub fn full( &self, matrix: E4m3BlockMatrix<'_>, activations: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Sourcepub fn column_parallel(
&self,
matrix: E4m3BlockMatrix<'_>,
activations: &[f32],
tokens: usize,
) -> Result<ColumnParallelResult, Box<dyn Error>>
pub fn column_parallel( &self, matrix: E4m3BlockMatrix<'_>, activations: &[f32], tokens: usize, ) -> Result<ColumnParallelResult, Box<dyn Error>>
Column-parallel projection. Weight output rows and their scale rows are partitioned across ranks. The input is host-broadcast, rank-local projections execute independently, and the output is host-gathered in rank order.
pub fn upload_column_parallel( &self, matrix: E4m3BlockMatrix<'_>, ) -> Result<ResidentColumnParallel, Box<dyn Error>>
pub fn column_parallel_resident( &self, matrix: &ResidentColumnParallel, activations: &[f32], tokens: usize, ) -> Result<ColumnParallelResult, Box<dyn Error>>
Sourcepub fn row_parallel(
&self,
matrix: E4m3BlockMatrix<'_>,
activations: &[f32],
tokens: usize,
) -> Result<RowParallelResult, Box<dyn Error>>
pub fn row_parallel( &self, matrix: E4m3BlockMatrix<'_>, activations: &[f32], tokens: usize, ) -> Result<RowParallelResult, Box<dyn Error>>
Row-parallel projection. Weight/input columns and their scale columns are partitioned across ranks. Rank-local partials return through host memory and are reduced in stable rank order.
pub fn upload_row_parallel( &self, matrix: E4m3BlockMatrix<'_>, ) -> Result<ResidentRowParallel, Box<dyn Error>>
pub fn row_parallel_resident( &self, matrix: &ResidentRowParallel, activations: &[f32], tokens: usize, ) -> Result<RowParallelResult, Box<dyn Error>>
pub fn upload_bf16_column_parallel( &self, matrix: Bf16Matrix<'_>, ) -> Result<ResidentBf16ColumnParallel, Box<dyn Error>>
Sourcepub fn upload_step_bf16_column_parallel(
&self,
matrix: Bf16Matrix<'_>,
) -> Result<ResidentBf16ColumnParallel, Box<dyn Error>>
pub fn upload_step_bf16_column_parallel( &self, matrix: Bf16Matrix<'_>, ) -> Result<ResidentBf16ColumnParallel, Box<dyn Error>>
Step-3.7 column projection with one numerical program across TP1/TP2/TP4/TP8.
Sourcepub fn upload_step_bf16_column_parallel_f32_mirror(
&self,
matrix: Bf16Matrix<'_>,
) -> Result<ResidentBf16ColumnParallel, Box<dyn Error>>
pub fn upload_step_bf16_column_parallel_f32_mirror( &self, matrix: Bf16Matrix<'_>, ) -> Result<ResidentBf16ColumnParallel, Box<dyn Error>>
Load-time exact F32 expansion of a Step BF16 shard.
The original BF16 allocation is released after the stream-ordered conversion. Decode then reuses the resident F32 values with the same topology-invariant output-row chunks.
pub fn bf16_column_parallel_resident( &self, matrix: &ResidentBf16ColumnParallel, activations: &[f32], tokens: usize, ) -> Result<ColumnParallelResult, Box<dyn Error>>
Sourcepub fn bf16_column_parallel_resident_native(
&self,
matrix: &ResidentBf16ColumnParallel,
activations: &[f32],
tokens: usize,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn bf16_column_parallel_resident_native( &self, matrix: &ResidentBf16ColumnParallel, activations: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Native-P2P twin of Self::bf16_column_parallel_resident.
The host-canonical activation is uploaded once on rank zero and peer-broadcast to the remaining ranks. Rank-local outputs are peer-gathered in token-major order before one root readback. This removes per-rank host staging but deliberately still returns a host oracle; attention and KV ownership are separate milestones.
Does the serving engine live in the SAME CUDA context as this runtime’s root rank?
The device-resident input/output seams below hand raw device buffers across the
Engine boundary, which is only addressable when both sides share the root device’s
primary context — the seam step35_tp_qkv keys its residency dispatch on.
Sourcepub fn bf16_column_parallel_resident_native_device(
&self,
matrix: &ResidentBf16ColumnParallel,
root_activation: &CudaSlice<f32>,
tokens: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn bf16_column_parallel_resident_native_device( &self, matrix: &ResidentBf16ColumnParallel, root_activation: &CudaSlice<f32>, tokens: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Device-input twin of Self::bf16_column_parallel_resident_native (lane/
hermes-perf-fixes, 2026-08-23 — the step QKV TP host-bounce finding). The activation
arrives as a ROOT-DEVICE buffer (first tokens * in_features values) instead of a
host slice, and the gathered output stays root-resident: no DtoH of the hidden state,
no host q/k/v staging, no re-upload. BYTE-IDENTICAL to the host-canonical native arm
by construction — the root input bytes are dtod-copied where the host arm htod’d the
same bytes, and every kernel, peer copy, and gather order is shared.
FENCES: caller must have synchronized the producer stream that wrote
root_activation (the serving engine’s — a DIFFERENT stream in the same context);
this method synchronizes the root stream before returning so the caller’s stream can
consume the gathered output immediately.
Sourcepub fn bf16_column_parallel_resident_device_shards_from_root(
&self,
matrix: &ResidentBf16ColumnParallel,
root_activation: &CudaSlice<f32>,
tokens: usize,
) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
pub fn bf16_column_parallel_resident_device_shards_from_root( &self, matrix: &ResidentBf16ColumnParallel, root_activation: &CudaSlice<f32>, tokens: usize, ) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
Root-device-input twin of Self::bf16_column_parallel_resident_device_shards:
the canonical activation is already resident on the root device (len >=
tokens * in_features; extra tail values beyond the active prefix are ignored,
the reused-prime-slab contract of active_matrix_values).
Sourcepub fn bf16_column_parallel_resident_device_shards(
&self,
matrix: &ResidentBf16ColumnParallel,
activations: &[f32],
tokens: usize,
) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
pub fn bf16_column_parallel_resident_device_shards( &self, matrix: &ResidentBf16ColumnParallel, activations: &[f32], tokens: usize, ) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
Keep Step BF16 column outputs resident on their owning TP ranks.
Rank zero receives the host-canonical activation once and peer-broadcasts it when TP>1.
Unlike Self::bf16_column_parallel_resident_native, this method performs no output
gather or readback. It is the correctness substrate for rank-local norm, RoPE, attention,
and cache ownership; callers must not treat its existence as serving qualification.
Sourcepub fn allocate_replicated_device_rows(
&self,
tokens: usize,
width: usize,
) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
pub fn allocate_replicated_device_rows( &self, tokens: usize, width: usize, ) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
Allocate one fixed-shape replicated batch without initializing its contents.
Callers must refresh every rank before passing the batch to an operator.
Sourcepub fn refresh_replicated_device_rows_from_root(
&self,
rows: &mut ResidentReplicatedDeviceRows,
source: &CudaSlice<f32>,
) -> Result<(), Box<dyn Error>>
pub fn refresh_replicated_device_rows_from_root( &self, rows: &mut ResidentReplicatedDeviceRows, source: &CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>
Replace a fixed-shape replicated batch from a root-device source.
Sourcepub fn upload_replicated_device_rows(
&self,
rows: &[f32],
tokens: usize,
width: usize,
) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
pub fn upload_replicated_device_rows( &self, rows: &[f32], tokens: usize, width: usize, ) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
Upload one canonical batch on rank zero and replicate it over native P2P.
Sourcepub fn bf16_column_parallel_resident_replicated_device_shards(
&self,
matrix: &ResidentBf16ColumnParallel,
activations: &ResidentReplicatedDeviceRows,
) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
pub fn bf16_column_parallel_resident_replicated_device_shards( &self, matrix: &ResidentBf16ColumnParallel, activations: &ResidentReplicatedDeviceRows, ) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
Execute a column-parallel BF16 matrix directly from rank-local replicated inputs.
Sourcepub fn upload_sigmoid_topk_router(
&self,
weight: Bf16Matrix<'_>,
correction_bias: &[f32],
active: Option<&[bool]>,
experts_per_token: usize,
scaling_factor: f32,
route_norm: bool,
) -> Result<ResidentSigmoidTopKRouter, Box<dyn Error>>
pub fn upload_sigmoid_topk_router( &self, weight: Bf16Matrix<'_>, correction_bias: &[f32], active: Option<&[bool]>, experts_per_token: usize, scaling_factor: f32, route_norm: bool, ) -> Result<ResidentSigmoidTopKRouter, Box<dyn Error>>
Upload a BF16 router once on rank zero and retain its exact F32 expansion.
Sourcepub fn sigmoid_topk_replicated_device_rows_host(
&self,
router: &ResidentSigmoidTopKRouter,
input: &ResidentReplicatedDeviceRows,
) -> Result<SigmoidTopKHostOutput, Box<dyn Error>>
pub fn sigmoid_topk_replicated_device_rows_host( &self, router: &ResidentSigmoidTopKRouter, input: &ResidentReplicatedDeviceRows, ) -> Result<SigmoidTopKHostOutput, Box<dyn Error>>
Route rank-zero replicated rows and return the narrow host control result plus logits.
The logits readback exists for independent oracle comparison. This method is a correctness surface; a serving scheduler may retain logits and selected routes on device.
Sourcepub fn upload_replicated_bf16_swiglu(
&self,
gate: Bf16Matrix<'_>,
up: Bf16Matrix<'_>,
down: Bf16Matrix<'_>,
) -> Result<ResidentReplicatedBf16SwiGlu, Box<dyn Error>>
pub fn upload_replicated_bf16_swiglu( &self, gate: Bf16Matrix<'_>, up: Bf16Matrix<'_>, down: Bf16Matrix<'_>, ) -> Result<ResidentReplicatedBf16SwiGlu, Box<dyn Error>>
Replicate a full BF16 SwiGLU bank on every rank.
Sourcepub fn replicated_bf16_swiglu_resident_device(
&self,
mlp: &ResidentReplicatedBf16SwiGlu,
input: &ResidentReplicatedDeviceRows,
activation_limit: Option<f32>,
) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
pub fn replicated_bf16_swiglu_resident_device( &self, mlp: &ResidentReplicatedBf16SwiGlu, input: &ResidentReplicatedDeviceRows, activation_limit: Option<f32>, ) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
Execute a fully replicated BF16 SwiGLU directly from replicated device rows.
Sourcepub fn rms_norm_replicated_device_rows(
&self,
input: &ResidentReplicatedDeviceRows,
weight: &[f32],
eps: f32,
) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
pub fn rms_norm_replicated_device_rows( &self, input: &ResidentReplicatedDeviceRows, weight: &[f32], eps: f32, ) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
Apply the same RMS-norm row program independently on every replicated rank.
Sourcepub fn add_rms_norm_replicated_device_rows(
&self,
input: &ResidentReplicatedDeviceRows,
update: &ResidentReplicatedDeviceRows,
weight: &[f32],
eps: f32,
) -> Result<(ResidentReplicatedDeviceRows, ResidentReplicatedDeviceRows), Box<dyn Error>>
pub fn add_rms_norm_replicated_device_rows( &self, input: &ResidentReplicatedDeviceRows, update: &ResidentReplicatedDeviceRows, weight: &[f32], eps: f32, ) -> Result<(ResidentReplicatedDeviceRows, ResidentReplicatedDeviceRows), Box<dyn Error>>
Add two replicated batches and RMS-normalize the exact residual on every rank.
pub fn collect_replicated_device_rows( &self, rows: &ResidentReplicatedDeviceRows, ) -> Result<Vec<Vec<f32>>, Box<dyn Error>>
pub fn upload_bf16_row_parallel( &self, matrix: Bf16Matrix<'_>, ) -> Result<ResidentBf16RowParallel, Box<dyn Error>>
pub fn bf16_row_parallel_resident( &self, matrix: &ResidentBf16RowParallel, activations: &[f32], tokens: usize, ) -> Result<RowParallelResult, Box<dyn Error>>
Sourcepub fn upload_step_bf16_row_parallel(
&self,
matrix: Bf16Matrix<'_>,
) -> Result<ResidentStepBf16RowParallel, Box<dyn Error>>
pub fn upload_step_bf16_row_parallel( &self, matrix: Bf16Matrix<'_>, ) -> Result<ResidentStepBf16RowParallel, Box<dyn Error>>
Step-3.7 row projection split into the same eight global K blocks for TP1/TP2/TP4/TP8.
pub fn upload_step_bf16_row_parallel_f32_mirror( &self, matrix: Bf16Matrix<'_>, ) -> Result<ResidentStepBf16RowParallel, Box<dyn Error>>
Sourcepub fn step_bf16_row_parallel_resident(
&self,
matrix: &ResidentStepBf16RowParallel,
activations: &[f32],
tokens: usize,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn step_bf16_row_parallel_resident( &self, matrix: &ResidentStepBf16RowParallel, activations: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Host-staged exactness twin of Self::step_bf16_row_parallel_resident_native.
Block inputs and partials cross host memory, but every partial is added on the root device in global checkpoint-column order. Native transport must reproduce this result bitwise.
Sourcepub fn step_bf16_row_parallel_resident_native(
&self,
matrix: &ResidentStepBf16RowParallel,
activations: &[f32],
tokens: usize,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn step_bf16_row_parallel_resident_native( &self, matrix: &ResidentStepBf16RowParallel, activations: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Native-P2P Step row projection with canonical global K-block reduction.
The full activation is uploaded once on the root. Each TP8-sized block is peer-scattered to its owning rank, its BF16 partial is peer-returned to the root, and root-device adds replay the same eight-block order as TP1 and the host-staged oracle.
Sourcepub fn step_bf16_row_parallel_resident_native_device(
&self,
matrix: &ResidentStepBf16RowParallel,
root_activation: &CudaSlice<f32>,
tokens: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn step_bf16_row_parallel_resident_native_device( &self, matrix: &ResidentStepBf16RowParallel, root_activation: &CudaSlice<f32>, tokens: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Device-input twin of Self::step_bf16_row_parallel_resident_native (lane/
hermes-perf-fixes, 2026-08-23): the full activation arrives as a ROOT-DEVICE buffer
and the reduced output stays root-resident — no DtoH of the attention output, no
host O staging, no re-upload. Byte-identical to the host-canonical arm by
construction (same block scatter, kernels, and global TP8 reduction order; the root
bytes are dtod-copied where the host arm htod’d the same bytes). Caller must have
synchronized the producer stream; the root stream is synchronized before returning.
Sourcepub fn step_bf16_row_parallel_resident_root_device(
&self,
matrix: &ResidentStepBf16RowParallel,
rank_activations: &[CudaSlice<f32>],
tokens: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn step_bf16_row_parallel_resident_root_device( &self, matrix: &ResidentStepBf16RowParallel, rank_activations: &[CudaSlice<f32>], tokens: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Reduce rank-local Step attention shards in canonical TP8 K-block order and keep the result on the root device.
Sourcepub fn step_bf16_row_parallel_resident_replicated_device(
&self,
matrix: &ResidentStepBf16RowParallel,
rank_activations: &[CudaSlice<f32>],
tokens: usize,
) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
pub fn step_bf16_row_parallel_resident_replicated_device( &self, matrix: &ResidentStepBf16RowParallel, rank_activations: &[CudaSlice<f32>], tokens: usize, ) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
Reduce rank-local Step attention shards, then replicate the canonical root result.
pub fn upload_expert( &self, gate: E4m3BlockMatrix<'_>, up: E4m3BlockMatrix<'_>, down: E4m3BlockMatrix<'_>, ) -> Result<ResidentTpExpert, Box<dyn Error>>
pub fn run_expert( &self, expert: &ResidentTpExpert, input: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
pub fn upload_expert_parallel( &self, gate: E4m3ExpertBank<'_>, up: E4m3ExpertBank<'_>, down: E4m3ExpertBank<'_>, ) -> Result<ResidentExpertParallel, Box<dyn Error>>
Sourcepub fn prepare_step_grouped_fp8_gate(
&self,
gate: E4m3ExpertBank<'_>,
up: E4m3ExpertBank<'_>,
down: E4m3ExpertBank<'_>,
input: &[f32],
tokens: usize,
selected: &[usize],
activation_limit: Option<f32>,
) -> Result<PreparedStepGroupedFp8Gate, Box<dyn Error>>
pub fn prepare_step_grouped_fp8_gate( &self, gate: E4m3ExpertBank<'_>, up: E4m3ExpertBank<'_>, down: E4m3ExpertBank<'_>, input: &[f32], tokens: usize, selected: &[usize], activation_limit: Option<f32>, ) -> Result<PreparedStepGroupedFp8Gate, Box<dyn Error>>
Prepare the official Step gate-only grouped-FP8 projection oracle on rank zero.
This intentionally does not alter the resident EP path. It owns a full rank-local tensor bank solely so the grouped projection can be compared with the existing per-route oracle without routing, transport, or combine changing underneath it.
Sourcepub fn run_step_grouped_fp8_gate(
&self,
plan: &mut PreparedStepGroupedFp8Gate,
) -> Result<StepGroupedFp8ProjectionOutput, Box<dyn Error>>
pub fn run_step_grouped_fp8_gate( &self, plan: &mut PreparedStepGroupedFp8Gate, ) -> Result<StepGroupedFp8ProjectionOutput, Box<dyn Error>>
Execute one prepared gate/up/activation/down projection sequence on rank zero.
pub fn prepare_step_grouped_expert_parallel_gate( &self, experts: &ResidentExpertParallel, input: &[f32], tokens: usize, selected: &[usize], activation_limit: Option<f32>, ) -> Result<PreparedStepGroupedExpertParallelGate, Box<dyn Error>>
pub fn prepare_step_grouped_expert_parallel_gate_with_capacity( &self, experts: &ResidentExpertParallel, input: &[f32], tokens: usize, selected: &[usize], activation_limit: Option<f32>, max_tokens: usize, ) -> Result<PreparedStepGroupedExpertParallelGate, Box<dyn Error>>
pub fn refresh_step_grouped_expert_parallel_gate( &self, experts: &ResidentExpertParallel, plan: &mut PreparedStepGroupedExpertParallelGate, input: &[f32], tokens: usize, selected: &[usize], ) -> Result<(), Box<dyn Error>>
Sourcepub fn refresh_step_grouped_expert_parallel_gate_from_root_device(
&self,
experts: &ResidentExpertParallel,
plan: &mut PreparedStepGroupedExpertParallelGate,
input: &CudaSlice<f32>,
tokens: usize,
selected: &[usize],
) -> Result<(), Box<dyn Error>>
pub fn refresh_step_grouped_expert_parallel_gate_from_root_device( &self, experts: &ResidentExpertParallel, plan: &mut PreparedStepGroupedExpertParallelGate, input: &CudaSlice<f32>, tokens: usize, selected: &[usize], ) -> Result<(), Box<dyn Error>>
Refresh routes and inputs from an already-resident rank-zero activation.
The caller must order the source producer before this call. The root copy is completed before peer dispatch, while CSR and workspace allocations retain their stable addresses.
Sourcepub fn refresh_step_grouped_expert_parallel_inputs_from_replicated(
&self,
experts: &ResidentExpertParallel,
plan: &mut PreparedStepGroupedExpertParallelGate,
input: &ResidentReplicatedDeviceRows,
) -> Result<(), Box<dyn Error>>
pub fn refresh_step_grouped_expert_parallel_inputs_from_replicated( &self, experts: &ResidentExpertParallel, plan: &mut PreparedStepGroupedExpertParallelGate, input: &ResidentReplicatedDeviceRows, ) -> Result<(), Box<dyn Error>>
Replace a fixed route plan’s rank inputs from an already replicated device batch.
Route CSR remains unchanged. Advancing the generation invalidates every prior projection and combine result, so callers must refresh combine metadata before executing again.
pub fn execute_step_grouped_expert_parallel_gate( &self, experts: &ResidentExpertParallel, plan: &mut PreparedStepGroupedExpertParallelGate, ) -> Result<(), Box<dyn Error>>
pub fn collect_step_grouped_expert_parallel_gate( &self, plan: &PreparedStepGroupedExpertParallelGate, ) -> Result<StepGroupedFp8ProjectionOutput, Box<dyn Error>>
pub fn run_step_grouped_expert_parallel_gate( &self, experts: &ResidentExpertParallel, plan: &mut PreparedStepGroupedExpertParallelGate, ) -> Result<StepGroupedFp8ProjectionOutput, Box<dyn Error>>
pub fn prepare_step_grouped_expert_parallel_combine( &self, plan: &PreparedStepGroupedExpertParallelGate, route_weights: &[f32], ) -> Result<PreparedPeerWeightedRouteCombine, Box<dyn Error>>
pub fn refresh_step_grouped_expert_parallel_combine( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &mut PreparedPeerWeightedRouteCombine, route_weights: &[f32], ) -> Result<(), Box<dyn Error>>
pub fn execute_step_grouped_expert_parallel_combine( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &mut PreparedPeerWeightedRouteCombine, ) -> Result<(), Box<dyn Error>>
pub fn collect_step_grouped_expert_parallel_combine( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &PreparedPeerWeightedRouteCombine, ) -> Result<Vec<f32>, Box<dyn Error>>
Sourcepub fn copy_step_grouped_expert_parallel_combine_root(
&self,
plan: &PreparedStepGroupedExpertParallelGate,
combine: &PreparedPeerWeightedRouteCombine,
destination: &Engine,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn copy_step_grouped_expert_parallel_combine_root( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &PreparedPeerWeightedRouteCombine, destination: &Engine, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Copy the active root combine result into a caller-owned engine on the same CUDA device.
The persistent combine buffer remains reusable by the next route generation; the returned allocation follows the serving runtime’s ordinary transient-output ownership.
pub fn broadcast_step_grouped_expert_parallel_combine( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &mut PreparedPeerWeightedRouteCombine, ) -> Result<(), Box<dyn Error>>
pub fn collect_step_grouped_expert_parallel_broadcast( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &PreparedPeerWeightedRouteCombine, ) -> Result<Vec<Vec<f32>>, Box<dyn Error>>
Sourcepub fn finish_step_grouped_expert_parallel_layer(
&self,
plan: &PreparedStepGroupedExpertParallelGate,
combine: &PreparedPeerWeightedRouteCombine,
shared: &ResidentReplicatedDeviceRows,
residual: &ResidentReplicatedDeviceRows,
) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
pub fn finish_step_grouped_expert_parallel_layer( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &PreparedPeerWeightedRouteCombine, shared: &ResidentReplicatedDeviceRows, residual: &ResidentReplicatedDeviceRows, ) -> Result<ResidentReplicatedDeviceRows, Box<dyn Error>>
Add routed and replicated shared-expert outputs, then add the attention residual.
pub fn run_step_grouped_expert_parallel_combine( &self, plan: &PreparedStepGroupedExpertParallelGate, combine: &mut PreparedPeerWeightedRouteCombine, ) -> Result<Vec<f32>, Box<dyn Error>>
pub fn upload_tensor_parallel( &self, gate: E4m3ExpertBank<'_>, up: E4m3ExpertBank<'_>, down: E4m3ExpertBank<'_>, ) -> Result<ResidentTensorParallel, Box<dyn Error>>
pub fn run_tensor_parallel_routes( &self, experts: &ResidentTensorParallel, input: &[f32], tokens: usize, selected: &[usize], route_weights: &[f32], experts_per_token: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Sourcepub fn gather_native_column_shards_device(
&self,
shards: &[CudaSlice<f32>],
tokens: usize,
local_out: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn gather_native_column_shards_device( &self, shards: &[CudaSlice<f32>], tokens: usize, local_out: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Gather token-major rank-local columns into one canonical root-device matrix.
pub fn gather_native_column_shards( &self, shards: &[CudaSlice<f32>], tokens: usize, local_out: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Sourcepub fn decode_v2_input_qkv_tcol(
&self,
ws_index: usize,
e: &Engine,
h_t: &CudaSlice<f32>,
t: usize,
q_m: &ResidentBf16ColumnParallel,
k_m: &ResidentBf16ColumnParallel,
v_m: &ResidentBf16ColumnParallel,
gate_shards: Option<StepTpGateShards<'_>>,
) -> Result<(), Box<dyn Error>>
pub fn decode_v2_input_qkv_tcol( &self, ws_index: usize, e: &Engine, h_t: &CudaSlice<f32>, t: usize, q_m: &ResidentBf16ColumnParallel, k_m: &ResidentBf16ColumnParallel, v_m: &ResidentBf16ColumnParallel, gate_shards: Option<StepTpGateShards<'_>>, ) -> Result<(), Box<dyn Error>>
v2 phase 1: replicate the layer input, project QKV, norm, rope, and stage the gate — all into the persistent workspace, ordered by events instead of host syncs.
The caller must have queued every producer of h, pos_d, and gate_raw on e’s
stream BEFORE this call: ev_entry is recorded once here and every rank stream waits
on it (the entry fence also guards workspace reuse across layers — any consumer of the
previous layer’s outputs was queued on e’s stream before this record).
T-COLUMN verify precompute (spec MTP): stage T input rows to every rank and run the
weight-amortized qkvg_tcol per rank into the ws slabs. Rope/norm/append stay per
column in the unmodified t=1 program (defer_norm_rope contract). Bit-exact per
column vs the t=1 kernel by construction.
pub fn run_routed_experts( &self, experts: &ResidentExpertParallel, input: &[f32], tokens: usize, selected: &[usize], route_weights: &[f32], experts_per_token: usize, activation_limit: Option<f32>, ) -> Result<Vec<f32>, Box<dyn Error>>
Source§impl TpE4m3HostBounce
impl TpE4m3HostBounce
Sourcepub fn full_nvfp4(
&self,
matrix: Nvfp4BlockMatrix<'_>,
activations: &[f32],
tokens: usize,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn full_nvfp4( &self, matrix: Nvfp4BlockMatrix<'_>, activations: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Unsharded NVFP4 projection on rank 0 (compatibility oracle). Macro applied post-kernel.
Sourcepub fn column_parallel_nvfp4(
&self,
matrix: Nvfp4BlockMatrix<'_>,
activations: &[f32],
tokens: usize,
) -> Result<ColumnParallelResult, Box<dyn Error>>
pub fn column_parallel_nvfp4( &self, matrix: Nvfp4BlockMatrix<'_>, activations: &[f32], tokens: usize, ) -> Result<ColumnParallelResult, Box<dyn Error>>
Column-parallel NVFP4 projection: output rows partition across ranks, host gather in rank order, macro applied ONCE post-gather.
Sourcepub fn row_parallel_nvfp4(
&self,
matrix: Nvfp4BlockMatrix<'_>,
activations: &[f32],
tokens: usize,
) -> Result<RowParallelResult, Box<dyn Error>>
pub fn row_parallel_nvfp4( &self, matrix: Nvfp4BlockMatrix<'_>, activations: &[f32], tokens: usize, ) -> Result<RowParallelResult, Box<dyn Error>>
Row-parallel NVFP4 projection: input columns partition at 64-superblock boundaries, rank-local partials reduce in stable rank order, macro applied ONCE post-reduce.
pub fn upload_expert_nvfp4( &self, gate: Nvfp4BlockMatrix<'_>, up: Nvfp4BlockMatrix<'_>, down: Nvfp4BlockMatrix<'_>, ) -> Result<ResidentTpNvfp4Expert, Box<dyn Error>>
pub fn run_expert_nvfp4( &self, expert: &ResidentTpNvfp4Expert, input: &[f32], tokens: usize, ) -> Result<Vec<f32>, Box<dyn Error>>
Sourcepub fn upload_tensor_parallel_nvfp4(
&self,
gate: Nvfp4ExpertBank<'_>,
up: Nvfp4ExpertBank<'_>,
down: Nvfp4ExpertBank<'_>,
) -> Result<ResidentNvfp4TensorParallel, Box<dyn Error>>
pub fn upload_tensor_parallel_nvfp4( &self, gate: Nvfp4ExpertBank<'_>, up: Nvfp4ExpertBank<'_>, down: Nvfp4ExpertBank<'_>, ) -> Result<ResidentNvfp4TensorParallel, Box<dyn Error>>
Upload every expert’s TP shards resident (one repacked block buffer per expert per rank).
Sourcepub fn upload_expert_parallel_nvfp4(
&self,
gate: Nvfp4ExpertBank<'_>,
up: Nvfp4ExpertBank<'_>,
down: Nvfp4ExpertBank<'_>,
) -> Result<ResidentNvfp4ExpertParallel, Box<dyn Error>>
pub fn upload_expert_parallel_nvfp4( &self, gate: Nvfp4ExpertBank<'_>, up: Nvfp4ExpertBank<'_>, down: Nvfp4ExpertBank<'_>, ) -> Result<ResidentNvfp4ExpertParallel, Box<dyn Error>>
Upload whole experts per owning rank (NVFP4 expert-parallel: the layout the clamped tail layers require — clamp semantics do not distribute across a tensor shard). Each owned expert keeps its full gate/up/down as one repacked block buffer on its owner.
Sourcepub fn run_routed_experts_nvfp4(
&self,
experts: &ResidentNvfp4ExpertParallel,
input: &[f32],
tokens: usize,
selected: &[usize],
route_weights: &[f32],
experts_per_token: usize,
activation_limit: Option<f32>,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn run_routed_experts_nvfp4( &self, experts: &ResidentNvfp4ExpertParallel, input: &[f32], tokens: usize, selected: &[usize], route_weights: &[f32], experts_per_token: usize, activation_limit: Option<f32>, ) -> Result<Vec<f32>, Box<dyn Error>>
Routed NVFP4 expert-parallel program, host-canonical: every selected expert executes WHOLE
on its owning rank (gate -> up -> clamped-or-plain SwiGLU on host -> down), each projection
macro applied once post-kernel, route-weighted accumulate on the host in slot order. The
activation uses step_expert_activation_host, so the clamped tail layers keep the official
contract. Exactness-first; no throughput claim.
Sourcepub fn run_tensor_parallel_routes_nvfp4_device(
&self,
experts: &ResidentNvfp4TensorParallel,
input: &[f32],
selected: &[usize],
route_weights: &[f32],
experts_per_token: usize,
activation_limit: Option<f32>,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn run_tensor_parallel_routes_nvfp4_device( &self, experts: &ResidentNvfp4TensorParallel, input: &[f32], selected: &[usize], route_weights: &[f32], experts_per_token: usize, activation_limit: Option<f32>, ) -> Result<Vec<f32>, Box<dyn Error>>
Device-resident routed NVFP4 expert program (decode shape, t=1 rows). The geometry gift this exploits: gate/up column halves land on the SAME rank that owns the matching down canonical shard (act[rank r] is exactly down-shard r’s input-column window), so the whole expert interior — gate, up, macro-scaled SwiGLU, down partial, route-weighted accumulate — runs rank-local with ZERO cross-rank transfer. Per (token, layer): one input upload per rank, one fenced peer copy of the remote accumulator, one root add, one readback.
Numeric class: device silu (silu_mul_scaled) with gate/up macros folded as gs/us and the down macro folded into the accumulate scalar (weight * macro_down — exact, both are per-expert constants). This matches the owning-stage MoE dev-path semantics, NOT the host-canonical program bit-for-bit; gate it with argmax + relative bounds against the host-canonical oracle, and with repeat determinism against itself. Clamped layers refuse (they stay on the EP program).
Sourcepub fn run_tensor_parallel_routes_nvfp4_device_io(
&self,
experts: &ResidentNvfp4TensorParallel,
e: &Engine,
input_dev: &CudaSlice<f32>,
selected: &[usize],
route_weights: &[f32],
experts_per_token: usize,
activation_limit: Option<f32>,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn run_tensor_parallel_routes_nvfp4_device_io( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, input_dev: &CudaSlice<f32>, selected: &[usize], route_weights: &[f32], experts_per_token: usize, activation_limit: Option<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Device-IO twin of run_tensor_parallel_routes_nvfp4_device: the layer input arrives as
a device row on the model engine e and the combined output returns as a fresh
e-context row — no host round-trip, no host stream sync. Ordering is evented (the v2
attention discipline): ev_entry is recorded on e’s stream AFTER the caller queued
the input’s producer; each rank waits it before its peer read; the root reduce waits
every rank’s done event; e waits the root’s done event before copying out. The
program bytes are identical to the host-IO twin — dtoh/htod and dtod preserve f32 bits.
Sourcepub fn nvfp4_routes_prestage(
&self,
experts: &ResidentNvfp4TensorParallel,
e: &Engine,
input_dev: &CudaSlice<f32>,
) -> Result<bool, Box<dyn Error>>
pub fn nvfp4_routes_prestage( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, input_dev: &CudaSlice<f32>, ) -> Result<bool, Box<dyn Error>>
Device-routed twin of run_tensor_parallel_routes_nvfp4_device_io: the selection and
route weights arrive as the device router’s e-context outputs — the per-layer host
logits readback disappears. The fresh router outputs are staged into persistent
e-context buffers on e’s stream (never-free discipline) before the entry event; each
rank peer-reads them behind it. The down-macro fold happens in-kernel.
Prestage the routed-expert input: pull the shared row to every rank and quantize it
there, WITHOUT the selection — callable before the router so the rank chains overlap
it. No-op (returns false) when the workspace is not built yet or the door is off;
the routed run then does its own staging as before.
Sourcepub fn nvfp4_routes_prestage_with(
&self,
experts: &ResidentNvfp4TensorParallel,
e: &Engine,
input_dev: &CudaSlice<f32>,
rank1_router: impl FnOnce(&Engine, &CudaSlice<f32>, &mut CudaSlice<i32>, &mut CudaSlice<f32>) -> Result<bool, Box<dyn Error>>,
) -> Result<bool, Box<dyn Error>>
pub fn nvfp4_routes_prestage_with( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, input_dev: &CudaSlice<f32>, rank1_router: impl FnOnce(&Engine, &CudaSlice<f32>, &mut CudaSlice<i32>, &mut CudaSlice<f32>) -> Result<bool, Box<dyn Error>>, ) -> Result<bool, Box<dyn Error>>
nvfp4_routes_prestage with a PEER-ROUTER hook: after rank1’s input pull +
quantize, the hook may compute rank1’s route selection LOCALLY (replicated router —
deterministic kernels on identical input bits produce identical sel/w, so the
selection is bit-equal to the root’s). Returns true when it wrote sel/route_w; the
routed run then skips rank1’s sel pull.
Sourcepub fn run_tensor_parallel_routes_nvfp4_device_routed_tn(
&self,
experts: &ResidentNvfp4TensorParallel,
e: &Engine,
z_t: &CudaSlice<f32>,
sel_d: &CudaSlice<i32>,
w_d: &CudaSlice<f32>,
t: usize,
n_sel_col: usize,
activation_limit: Option<f32>,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn run_tensor_parallel_routes_nvfp4_device_routed_tn( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, z_t: &CudaSlice<f32>, sel_d: &CudaSlice<i32>, w_d: &CudaSlice<f32>, t: usize, n_sel_col: usize, activation_limit: Option<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
TWO-COLUMN device-routed expert program (spec verify, MEMRA_TCOL_FFN): one gu_tcol sweep over 2*n_sel_col pairs (pair t reads activation row t/n_sel_col — weights the two columns share dedup through L2), the UNCHANGED silu/down kernels at n_sel=16 (both already index per pair), and one offset-axpy combine per column (the exact t=1 sequential chain over that column’s 8 pairs). No serving doors: no graph, no prestage, no shexp folding — plain evented ordering. Returns [2, input_width] on e.
EXACTNESS: every kernel body is the t=1 program per (pair,row) or per element; the per-column combine order equals the t=1 combine; the cross-rank join adds the same operand values elementwise. Gated by the greedy tape like every verify arm.
pub fn run_tensor_parallel_routes_nvfp4_device_routed( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, input_dev: &CudaSlice<f32>, sel_d: &CudaSlice<i32>, w_d: &CudaSlice<f32>, experts_per_token: usize, activation_limit: Option<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Sourcepub fn run_tensor_parallel_routes_nvfp4_device_routed_prejoin(
&self,
experts: &ResidentNvfp4TensorParallel,
e: &Engine,
input_dev: &CudaSlice<f32>,
sel_d: &CudaSlice<i32>,
w_d: &CudaSlice<f32>,
experts_per_token: usize,
activation_limit: Option<f32>,
pre_join: impl FnOnce() -> Result<(), Box<dyn Error>>,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn run_tensor_parallel_routes_nvfp4_device_routed_prejoin( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, input_dev: &CudaSlice<f32>, sel_d: &CudaSlice<i32>, w_d: &CudaSlice<f32>, experts_per_token: usize, activation_limit: Option<f32>, pre_join: impl FnOnce() -> Result<(), Box<dyn Error>>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
run_tensor_parallel_routes_nvfp4_device_routed with a PREJOIN hook: pre_join
runs on the host right before the join wait is enqueued on e’s stream — work it
issues there (e.g. the shexp overlap) executes WHILE the peer rank finishes its
sweep, instead of after the join. Value-neutral by construction (the hook only
reorders independent host issue).
Sourcepub fn run_tensor_parallel_routes_nvfp4_device_routed_prejoin_add3(
&self,
experts: &ResidentNvfp4TensorParallel,
e: &Engine,
input_dev: &CudaSlice<f32>,
sel_d: &CudaSlice<i32>,
w_d: &CudaSlice<f32>,
experts_per_token: usize,
activation_limit: Option<f32>,
pre_join: impl FnOnce() -> Result<(), Box<dyn Error>>,
post_add: Option<(u64, u64)>,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn run_tensor_parallel_routes_nvfp4_device_routed_prejoin_add3( &self, experts: &ResidentNvfp4TensorParallel, e: &Engine, input_dev: &CudaSlice<f32>, sel_d: &CudaSlice<i32>, w_d: &CudaSlice<f32>, experts_per_token: usize, activation_limit: Option<f32>, pre_join: impl FnOnce() -> Result<(), Box<dyn Error>>, post_add: Option<(u64, u64)>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
The prejoin variant with MOE TAIL FUSION M1: when post_add = Some((sh_raw, scale_raw)), the direct-join arm folds the shexp apply into the join add
(dst = (acc0+acc1) + sh*scale[0], exact split-pair sequence) — the caller skips
its apply launch. Raw UVA pointers so no lock is held across the call.
Sourcepub fn run_tensor_parallel_routes_nvfp4(
&self,
experts: &ResidentNvfp4TensorParallel,
input: &[f32],
tokens: usize,
selected: &[usize],
route_weights: &[f32],
experts_per_token: usize,
activation_limit: Option<f32>,
) -> Result<Vec<f32>, Box<dyn Error>>
pub fn run_tensor_parallel_routes_nvfp4( &self, experts: &ResidentNvfp4TensorParallel, input: &[f32], tokens: usize, selected: &[usize], route_weights: &[f32], experts_per_token: usize, activation_limit: Option<f32>, ) -> Result<Vec<f32>, Box<dyn Error>>
Routed NVFP4 expert program, host-canonical transport. Native/bulk P2P transport for the NVFP4 bank is a separate increment; this entry point is exactness-first and reports no throughput claim.