Module tp 

Tensor-parallel correctness runtime.

This module is deliberately narrower than the serving runtime. It executes real rank-local E4M3 projections on distinct CUDA devices. Deterministic host-staged collectives remain the default exactness reference; an opt-in native-P2P path must reproduce the same canonical checkpoint-block program before it can advance. Neither path is product-throughput evidence.

Structs

Bf16Matrix

ColumnParallelResult

E4m3BlockMatrix

E4m3ExpertBank

Nvfp4BlockMatrix

One expert’s modelopt NVFP4 projection: packed codes + per-16 UE4M3 scale bytes + macro.

Nvfp4DeviceRoutesWorkspace

Nvfp4ExpertBank

Stacked modelopt NVFP4 expert bank (host view over the checkpoint bytes).

Nvfp4T2Workspace

Persistent buffers for the two-column (spec verify) NVFP4 device-routed program: every slab is the t=1 workspace shape doubled along the pair axis, plus per-column accumulators. One per expert bank, reused every (round, layer) call.

PreparedPeerWeightedRouteCombine

Persistent root-side weighted combine for peer-owned canonical route rows.

PreparedStepGroupedExpertParallelGate

Prepared official Step expert-owner grouped-FP8 projection gate.

PreparedStepGroupedFp8Gate

Prepared official Step grouped-FP8 projection gate.

ResidentBf16ColumnParallel

ResidentBf16RowParallel

ResidentColumnParallel

ResidentExpertParallel

Persistent expert-parallel reference.

ResidentNvfp4ColumnBankRank

One rank’s resident NVFP4 expert bank shard: one repacked block buffer PER expert (per-expert device allocations keep this increment off any new strided-kernel API; the strided twin is a later perf rung, mirroring the FP8 bank’s history).

ResidentNvfp4ColumnParallel

ResidentNvfp4ExpertParallel

ResidentNvfp4Rank

One rank’s resident repacked NVFP4 shard: memra block_nvfp4 rows on device.

ResidentNvfp4RowBankRank

ResidentNvfp4RowParallel

ResidentNvfp4TensorParallel

ResidentReplicatedBf16SwiGlu

Full BF16 SwiGLU weights replicated independently on every runtime rank.

ResidentReplicatedDeviceRows

One token-major F32 batch replicated across a native-P2P rank group.

ResidentRowParallel

ResidentSigmoidTopKRouter

Root-owned BF16 sigmoid router with persistent F32 weight, bias, and active mask.

ResidentStepBf16RowParallel

ResidentTensorParallel

Persistent tensor-parallel expert bank.

ResidentTpExpert

Persistent TP2/TP4/TP8 routed-expert reference.

ResidentTpKvCache

ResidentTpKvCacheRank

ResidentTpNvfp4Expert

RowParallelResult

SigmoidTopKHostOutput

StepEpLayerSpec

StepGroupedFp8ProjectionOutput

Projection-level output from the opt-in official Step grouped-FP8 gate.

StepTpDecodeV2Ws

TokenGraph

TokenGraphBuilder

TpE4m3HostBounce

Multi-context TP correctness runtime. Each rank owns an independent Engine and CUDA context.

TpKvAppendPlan

TpKvTransaction

Enums

KvRingAppend

StepTpGateShards

Persistent workspace of the v2 rank-local decode-attention driver.

Constants

NVFP4_CANONICAL_ROW_SHARDS

Canonical row-shard count for the NVFP4 down projection. The down reduction ALWAYS executes as exactly this many input-column windows summed in shard order, at every world size: a single full-width dot and a two-half-dots-plus-add differ in f32 parenthesization, so pinning the shard grid (not the world size) is what makes the TP1-oracle-vs-TP2 bit gate meaningful. This is the NVFP4 twin of the FP8 bank’s canonical checkpoint-block reduction.

Functions

graph_section

The section annotation: eager mode runs the closure verbatim; build mode wraps it in a stream capture on engine’s stream and records the child. Sections sharing a parallel_group id fork from the same predecessor set and merge together. The closure must be capture-safe (raw copies at cross-context seams, no host syncs, no events).

graph_section_absorbing

Serial section that ALSO joins every pending detached section (the SH1 consumer shape).

graph_section_detached

graph_section with the DETACHED shape: forks from the current frontier (or the open group base) and is joined only by the next serial section — never gates a group merge.

graph_section_opts

moe_residual_host

Canonical MoE output order: routed plus shared, then add the layer residual.

parse_step_ep_layer_specs

parse_step_tp_layer_specs

step_ep_device_arithmetic_enabled

step_ep_layer_specs

step_expert_activation_host

step_nvfp4_dev_routes_enabled

Opt-in door for the device-resident NVFP4 TP routed-expert decode program. Default OFF; the host-canonical program remains the oracle until the device path carries its own gates.

step_tp_bulk_p2p_enabled

step_tp_dcw_enabled

Device-counter attention path (graph increment A run EAGERLY): append at len_d - base_d, inc_i32, fa over the counter-derived window — with bucket = the effective t_kv this is bit-identical to the host-row + kvmod path (the one-partition law), and it is the exact child content the capture wraps. Rebase tokens and sub-vec-floor contexts fall back.

step_tp_decode_v2_enabled

The v2 rank-local Step decode-attention driver: persistent workspaces, evented cross-stream ordering, and a root-device O reduction — same kernels, values, and canonical reduction order as the v1 driver (it requires the F32 mirror so no per-call weight expansion exists on either side of the comparison).

step_tp_dev_router_enabled

Device-side sigmoid top-k routing for the TP device-IO expert program: the per-layer host logits readback (the last per-layer host sync) disappears. Selection tie-breaking may differ from the host router — NUMERIC-CLASS door, run-gen argmax gate + boot battery.

step_tp_f32_mirror_enabled

step_tp_graph_enabled

CUDA-graph door for the shape-stable TP segments (first increment: the device-routed expert program — per-layer multi-device parents built from per-rank children, launched on the model engine’s stream; zero per-token node updates). Mechanism proven by tp_graph_probe. VALUE-IDENTICAL: the graphs replay exactly the eager kernel/copy sequence.

step_tp_layer_specs

step_tp_native_p2p_enabled

step_tp_qkv_fused_enabled

Fused single-launch QKV projection inside the v2 decode driver — a NUMERIC-CLASS door (per-row deterministic tree reduce instead of the chunked cuBLASLt program), default OFF, gated by the run-gen argmax gate + boot battery like MEMRA_STEP_NVFP4_DEV_ROUTES.

token_graph_build_begin

Arm the thread-local builder (build mode) — the next graph_section calls capture.

token_graph_build_finish

Take the finished parent (ends build mode).

token_graph_building

True while the thread-local builder is armed.

Type Aliases

StepTpLayerSpec