Skip to content

Numerics envelopes ​

--numerics names, at launch level, the numerical contract a deployment promises. Every capability under it exists as an individual switch; the envelope's whole job is to keep the set coherent, because RL rollout introduced a class of deployment where one missing switch silently invalidates the training signal.

The contract ​

--numerics rl-bitwise is the one bitwise envelope. It promises, within one deployment (fixed world size, parallel layout, model, kernels):

  1. Run invariance — the same request produces bitwise-identical tokens and logprobs across runs.
  2. Batch invariance — a request's tokens and logprobs do not depend on which other requests share its batches, or on how the scheduler happened to chunk and batch it.
  3. Trainer alignment — the forward follows the training framework's operation order wherever the two engines are known to differ, so that a teacher-forced pass over a trainer-generated sequence reproduces the trainer's per-token logprobs bitwise. The switches this tightens are listed under alignment.trainer below; each exists individually for auto.

It deliberately does not promise (yet — the layer exists in the hierarchy, unimplemented):

  1. Topology invariance — the same logprobs under a different TP/DP factorization (needs TP-invariant projection layouts on top of the vocab-block log-softmax of alignment.trainer).

Two consequences of 3 are enforced at startup, so a model that cannot meet them is refused under rl-bitwise with the reason: MoE TP must be 1 (the slot-order MoE combine folds over the EP group inside the leaf), and the vocabulary must be a whole number of BLOCKWISE_VOCAB_BLOCK-wide blocks (the trainer's log-softmax folds fixed blocks).

The hierarchy ​

numerics.mode                       --numerics {auto, rl-bitwise}
├── kernels.deterministic           fixed-reduction-order compute
│   ├── no autotune                 disable_autotune (tactic choice is shape-
│   │                               and machine-dependent state); no
│   │                               persistent tactic cache is loaded
│   ├── no TF32                     disable_tf32 + NVIDIA_TF32_OVERRIDE=0
│   ├── no PDL                      disable_pdl (serialize kernel chains)
│   └── batch-invariant leaves      kernel registry: leaves declaring the
│                                   "batch_invariant" feature; a caller in
│                                   rl-bitwise REQUIRES the feature, so a
│                                   missing implementation fails at startup
│                                   instead of silently falling back
├── collectives.deterministic       one association order per reduction
│   ├── no fused AR+norm            enable_allreduce_fusion=False and
│   │                               comm_fusion_max_num_tokens=-1, also set
│   │                               by batch_invariant_collectives or
│   │                               layer_boundary_norm=unfused alone (the
│   │                               fused kernels make no bitwise claim)
│   ├── NCCL_ALGO=Ring,             the algorithm/protocol switch by message
│   │   NCCL_PROTO=Simple           size changes association order
│   ├── batch_invariant_collectives AutoBackend.route, one decision for every
│   │                               collective: every reduction is data
│   │                               movement plus a fixed-rank-order fp32
│   │                               fold (all-reduce = all-gather + fold at
│   │                               world_size x traffic, its gather on the
│   │                               multicast kernels where they reach, or
│   │                               trtllm workspace kernels with fp32
│   │                               accumulation where a workspace is armed
│   │                               (one call, or chunks of the rows it
│   │                               holds; armed through a prefill chunk
│   │                               up to 2048 rows): they fold ranks
│   │                               0..W-1 the same way on one-shot and
│   │                               two-shot, bit for bit, -0.0 sums
│   │                               included as +0.0;
│   │                               reduce-scatter = all-to-all + fold at no
│   │                               extra traffic).
│   │                               A ring reduction chunks by message size,
│   │                               so its per-element order is run-stable
│   │                               but not batch-size-invariant; the NVLS
│   │                               in-switch reduction's order follows the
│   │                               row's offset in the buffer (section
│   │                               below). MAX/MIN all-reduces are exact
│   │                               in any order and stay on NCCL. Gathers
│   │                               move data and keep the multicast
│   │                               kernels. The route reads only static
│   │                               properties of the call site, never the
│   │                               row count, so a site cannot change route
│   │                               with the batch. Inside the all-reduce
│   │                               fold the row count does pick the kernel
│   │                               (workspace one-shot or two-shot, or the
│   │                               gather where no armed workspace serves
│   │                               the shape); all give the same bits,
│   │                               except values below 1.2e-38, which the
│   │                               workspace kernels (built with fast
│   │                               math) flush to zero
│   ├── data movement               all-gather, token all-gather, all-to-all
│   │                               (even or uneven, the transposes of the
│   │                               decode TP layouts) move bytes without
│   │                               reducing: nothing to fold, no branch
│   └── force_deterministic_rsag    the user's "NCCL and the fold only" knob:
│                                   no symmetric-memory path at all (multicast
│                                   gathers, the trtllm and Iris all-reduce
│                                   tiers, distributed argmax). Not folded by
│                                   the envelope; under it the only change is
│                                   that gathers move to NCCL
├── sampling.deterministic          sampling_stream=per-request: sampled rows
│                                   draw from the Gumbel-max pool kernels,
│                                   whose stream is keyed by the request's
│                                   seed (explicit, else crc32(rid)) and the
│                                   absolute target token position. Prefill
│                                   adds its chunk length to the KV frontier;
│                                   decode adds one. Neither chunk counts nor
│                                   prefix hits enter the random key.
│                                   flashinfer's *_sampling_from_probs
│                                   kernels are NOT: they read one seed and
│                                   offset for the whole batch and seed
│                                   curand with the batch row, so a request's
│                                   draw moves with its co-batch (T>0
│                                   'packed' fails while 'rerun' passes).
│                                   N=1 verify uses the same pool sampler.
│                                   Multi-token verify keeps flashinfer's
│                                   per-slot coins; its cache/scheduling
│                                   invariance is not established here.
│                                   Greedy rows additionally take
│                                   the canonical lowest-index argmax, in
│                                   sampling and in speculative verify
│                                   (exact-match chain), because EXACT logit
│                                   ties happen in practice and the pool
│                                   route's stochastic kernels resolve them
│                                   in batch-shape-dependent reduction order;
│                                   backends without the overlay are refused.
│                                   Under --enable-speculative-sampling the
│                                   draft proposal is one more per-request
│                                   stream: Gumbel-max noise keyed by the
│                                   request's seed and a salted (position,
│                                   step) offset (never the batch row), the
│                                   verify coins stay per-slot, and greedy
│                                   rows keep the canonical argmax with a
│                                   one-hot q, so their verify is unchanged
├── invariance.batch                per-row-independent reductions
│   ├── no split-KV attention       decode kernels whose split count scales
│   │                               with batch/SM occupancy are excluded by
│   │                               the batch_invariant feature
│   ├── row-local top-k             the DSA indexer's selection resolves
│   │                               equal scores toward the lowest candidate
│   │                               within each row (dsa_*_topk
│   │                               batch_invariant=True); the tuned top-k
│   │                               kernels switch algorithm and CTA split
│   │                               with the row count, which moves ties
│   ├── selection-order reduction   dsa_slot_order=selection: reduce selected
│   │                               KV rows in the top-k leaf's emitted order,
│   │                               matching the trainer. Both selection and
│   │                               attention require batch-invariant leaves;
│   │                               physical cache addresses never determine
│   │                               the reduction order
│   └── per-row GEMMs               fixed-order GEMM leaves (see plans below)
├── logprob.topology-invariant      (deferred) TP-invariant projection
│                                   layouts on top of the vocab-block
│                                   log-softmax of alignment.trainer
└── alignment.trainer               the trainer's operation order (section
                                    below); rl-bitwise folds every switch

Precedence: the envelope only ever tightens. It sets every switch it governs to its tight value, and it refuses an explicit choice it cannot tighten — a named MoE backend other than the batch-invariant one, a sampling backend without canonical greedy ties — rather than keeping it and silently voiding the contract. resolve_numerics runs after resolve_communication so it can veto the auto-enabled all-reduce fusion. Every closed-set switch is validated once, in ServerArgs; the layers that read a resolved switch trust it, and a constraint that needs more than the launch (a model's vocabulary, its routing) is checked when that module is constructed, never per forward. Every selection point that pins a batch-invariant leaf tests numerics in BITWISE_ENVELOPES, never the one name, so an envelope added above rl-bitwise would inherit every pin.

Unquantized FP16/BF16 MoE on flashinfer_cutlass defaults to GEMM2 atomic expert combine. --no-moe-cutlass-fused-finalize selects a separate finalize reduction for both target and draft; rl-bitwise forces fusion off. The setting is fixed in each expert plan before warmup and CUDA graph capture.

alignment.trainer ​

--logprob-order blockwise-16384 and blockwise-direct-16384 select 16384-wide sum-exp blocks with the same subtraction forms as blockwise and blockwise-direct. rl-bitwise preserves each explicit choice. Select the trainer's vocabulary-shard width independently of serving TP: folding eight narrow partials is not the same FP32 arithmetic as folding four wide partials. Prompt and output logprobs share the selected operation; logits and sampled tokens are unchanged. The model's existing rl-bitwise vocabulary eligibility remains a multiple of 32768.

The invariance layers make a deployment agree with itself. Trainer alignment makes it agree with a different program: the training framework's forward, whose arithmetic was never written to match an inference engine. Each switch below replaces one operation the engine performs differently from the trainer with the trainer's form. The table says what each switch changes: forward values means every activation downstream moves (tokens can flip at ties, logprobs change); logprobs only means the sampled tokens are untouched and only the reported log-probabilities change.

SwitchTrainer formChanges
--sampling-stream per-requestNon-greedy rows draw from a Philox stream keyed by (request seed, position) only (sampling.deterministic above); the trainer plays back the sampled ids, so this is an invariance switch rather than an alignment one — T>0 rollouts need ittokens at T>0 (not logprobs)
--yarn-ramp-mask-device cpuCompute the deepseek_yarn linear ramp on CPU; rl-bitwise pins this choice. The computed ramp is transferred to the frequency device before interpolationforward values
--yarn-frequency-device cudaCompute YaRN powers, divisions and interpolation on CUDA after transferring the CPU-built ramp; rl-bitwise pins this choice, overriding an explicit cpu. CPU and CUDA arithmetic can round differentlyforward values
--mla-lora-scale runtimeThe sqrt(hidden / lora_rank) norm scales of LongCat-style MLA stay out of the q_a_layernorm / kv_a_layernorm weights and multiply q after q_b_proj and the latent after kv_a_layernorm in bf16, as the trainer does; the DSA indexer reads the unscaled q_loraforward values
--layer-boundary-norm unfusedThe norm that opens each physical layer and the final norm read a bf16 hidden + residual materialized first (residual = hidden), then a standalone RMSNorm, instead of the fused add+norm kernel whose sum stays fp32; all-reduce+norm fusion is vetoed with itforward values
--router-topk torchThe correction-bias router runs fp32 torch.softmax, torch.topk(probs + bias, sorted=True) (PyTorch tie order), weights = unbiased probs x routed_scaling_factor, zero experts (id >= num_real) become -1 and keep their weight for the identity residualforward values (expert selection at near-ties, weights)
--dsa-slot-order selectionKeep the top-k leaf's emitted order when reducing sparse attention, as the trainer does. rl-bitwise fixes this choice, overriding an explicit sorted; both top-k and attention leaves still require batch invariance. Under auto, sorted remains available to reduce by token position, never physical cache addressforward values
--logprob-order blockwiseSelected-token logprobs use blockwise reduction: row max, shift, target gather, sum_exp over fixed 32768-wide vocab blocks (the in-block tree is the registered sampling.block_sumexp leaf's; the fold across blocks is a rank-ordered fp32 left fold), negative cross-entropy (or direct subtraction with blockwise-direct), for output and prompt (input) logprobs alikelogprobs only
--moe-combine-order slotThe MoE leaf folds a token's top-k routed outputs in fp32 in slot order across the EP group itself and adds the identity zero-expert residual in the same fold, as the trainer's grouped MLP does (moe_plan(combine_order="slot"), plan["process_group"] = the EP group); the host hands it the raw top-k (zero-expert ids intact, weights kept), adds no residual and runs no MoE all-reduce / reduce-scatter: CommManager.post_moe_comm, the one MoE reduction point every model goes through, reads the switch and only takes back this rank's rows in the RSAG layout; all-reduce+norm fusion is vetoed (CommManager.should_fuse re-derives the veto, so it holds wherever the fused kernel is relied upon). Needs MoE TP 1 (refused by resolve_numerics otherwise) and a kernel declaring combine_order with slot; under --all2all-backend deepep that kernel must also own the exchange (supports_all_to_all_ep) and returns the complete output of the rows this rank owns. Under rank the kernel returns a per-rank partial, the host reduces rank by rank and LongCat's residual enters exactly one partial (adds_zero_expert_residual, tp_ep_rank 0)forward values

rl-bitwise folds YaRN to a CPU ramp and CUDA frequencies. The same construction can be selected under auto with --yarn-ramp-mask-device cpu --yarn-frequency-device cuda. This replaces the previous all-CPU inverse-frequency construction under rl-bitwise. Under auto, both devices default independently to cuda; setting only the ramp device to cpu leaves frequency construction on CUDA. For the previous all-CPU construction, set both options to cpu. Both devices participate in the RoPE cache identity and the L3 KV deployment identity, preventing reuse across construction devices. The transfer and frequency construction run when the cos/sin cache is initialized; forward reuses that cache.

These are the switches the host can mirror with no new vendor dependency. The rest of the trainer's form lives in kernel leaves and model code the host does not own and is the out-of-tree model's part of the bargain: a sampling.block_sumexp leaf with the trainer's in-block order, a MoE apply leaf declaring combine_order={"rank", "slot"} whose slot form applies the routing probabilities inside the activation and combines the zero-expert residual in fp32 slot order, BF16 index-K indexer scoring and top-k leaves, and the same LoRA-scale placement in the draft model.

The host's side of the indexer is the plane and the facades. A DSA model's configure-attention hook names its index-key storage on the model config (ModelConfig.index_k_format, read by DSAConfig): the in-tree configure_dsa_attention names fp8_scaled — FP8 keys plus per-128 fp32 scales, the in-tree leaves' plane — and a plugin hook that scores the checkpoint's keys unquantized names bf16; a hook that names none is a construction error. The hook receives the resolved launch (ModelProfile.configure_attention(model_config, server_args)), so a plugin can name the plane its leaves score under server_args.numerics and keep the FP8 plane under auto. The ordinary recipe plans that plane, the pool writes keys in the plane's own dtype and never converts between the two, and dsa_decode_topk / dsa_prefill_topk read index_k_format and index_k_layout off the plane's dtype and shape, so a bf16 plane selects only a leaf declaring index_k_format={"bf16"} (the kernel package's DSA README has the table); the GLM-5.3-Flash recipe plans pooled fp8_scaled rows and refuses any other plane. candidate_lens_cpu reaches every top-k leaf registered with the candidate_lens_cpu feature and no other. Index keys handed to dsa_prefill_topk as rows in workspace-row order (the query-context-parallel history gather over page-sharded caches, docs/design/unified_path.md) keep the plane's format -- index_k_fp8 + index_k_scale, or index_k_bf16 -- and reach only leaves declaring the index_k_workspace_rows feature for it (selection requires the feature, overrides included, and the keywords are routed by it), so a bf16 leaf that scores gathered rows declares that feature and takes the index_k_bf16 keyword; the GPU DSA leaf selects that leaf once at construction under query context parallelism, so a missing one fails at startup. In-tree drafts fold no LoRA norm scale; a draft that does must read --mla-lora-scale exactly as the target does, folding only under folded.

When a model owns a cache format or execution contract that differs from its architecture's default, its profile may provide create_attention_config(server_args, model_config, is_draft). The factory returns the model-wide AttnConfig and its softmax component, validates its supported dtype and topology, and uses the same resolved launch for target and draft construction. The runtime rejects a result whose is_draft differs from the requested role. It still adds the profile's linear-attention component through the ordinary composition path. Every profile must explicitly pass a factory or None; None keeps the architecture's existing generator. A factory error propagates instead of falling back to a different cache format or numerical path.

Trainer alignment is a stronger claim than invariance and cannot be checked by the engine alone: a model earns rl-bitwise in ModelProfile.numerics_envelopes through the invariance harness and the teacher-forced comparison against a trainer dump described under Acceptance. In-tree models do not declare it; the out-of-tree LongCat 2.0 plugin is the first candidate.

Why not the in-switch reduction ​

The NVLS multimem.ld_reduce sums the group's copies of a row in the switch (fp32 accumulate, one rounding to bf16) and moves the payload through one port twice, far under the fold's world_size x all-gather, so it was the batch-invariant all-reduce for a while: one fixed rank issued every row's loads, and a startup self-check verified on each group that an ill-conditioned payload reduced to the same bits across repetitions and that its first half reproduced the first half's rows. Both held, and neither covers what a batch changes. A token's row lands at a different offset in the buffer with every batch, and the switch's association order follows that offset: on 8xH20 at width 3072 or 6144, the same rows reduced one or more rows further into the buffer come back different on more than half of an ill-conditioned payload's elements, and occasionally on ordinary activations (one element of 1.5 million in the probe). The order also differs between issuing ranks and between GPU sets, so nothing in software can pin it.

One element in a million is enough. Under attention TP every layer all-reduces o_proj and the dense down_proj; a last-bit difference in one row changes the KV it writes, and every later token of that request inherits it. Measured on LongCat 2.0 at TP8 on GSM8K (128 greedy chat requests of up to 1024 tokens), the in-switch route left a quarter of the decode logprobs unequal to the prefill rescoring the same ids, 3-5% of the prompt logprobs unequal between two prefills of different batches, and tokens diverging between two machines; with the fold all of them were bitwise equal. Every batch-invariant reduction therefore takes the fold (AutoBackend.route), the in-switch kernels serve only the performance routes, and a deployment that cannot afford the fold's traffic removes the reduction instead: --tp-batch-invariant (below) leaves the decode layer no cross-rank sum outside MoE.

Parallel layouts and the envelope ​

The decode-side TP layouts under attention DP (--attn-head-tp-size, --lm-head-tp-size, --dense-tp-size, --tp-batch-invariant; see Parallelism) are explicit deployment choices. The envelope does not fold them in: a layout is part of the deployment the contract is stated for, not a switch the contract tightens, and rl-bitwise neither selects nor refuses one. Within one deployment they keep the contract as follows.

  • Run and batch invariance hold under every layout. The head, dense and LM-head GEMMs take the same batch-invariant plans as their replicated forms; the exchanges (token all-gather, all-gather of a reduction dimension, the even and uneven all-to-all transposes) are permutations of bytes and need no fold; the remaining reductions (o_proj and down_proj reduce-scatter without --tp-batch-invariant) go through batch_invariant_collectives.
  • --tp-batch-invariant is the stronger property: with it the decode layer has no cross-rank reduction outside MoE. o_proj and down_proj become column-parallel on hidden fed by an all-gather of the reduction dimension, so every output element is one full-K GEMM result — bitwise the value a TP1 or replicated layer computes, independent of the width W. An ordered-fold reduce-scatter is batch-invariant too, but its fp32 fold of W partials is not the full-K GEMM's association, so a decode engine on that path disagrees in the last bits with a prefill engine whose o_proj / down_proj run replicated or at TP1. resolve_numerics therefore warns when head TP runs under rl-bitwise without --tp-batch-invariant attn.
  • Topology invariance across layouts is not promised beyond that: the vocab-sharded log-softmax and the trainer-aligned fold remain the deferred layers above, and a deployment advertising rl-bitwise with one of these layouts must pass the invariance harness with that layout.

Layout invariance of query context parallelism ​

Every collective query context parallelism adds (docs/design/unified_path.md) is data movement: row selection, the all-gather of rotated latent rows, index-K rows, gathered history rows, sampled rows and the planned prompt-logprob rows' activations, and the permutations between gather order and span order. The zigzag split of every request over the ranks only decides which rank computes a row. The per-row kernels — sparse attention over the gathered history with every head and no LSE merge, the indexer's top-k over pre-gathered rows, RoPE, the GEMMs — see for each row exactly the operands a single GPU would, so the permutation itself is byte-preserving. Bitwise equivalence also requires the selected kernels and reductions to be batch-invariant; identical operands alone do not ensure identical rounding when row counts select different kernel tactics. The prompt logprobs in particular are the tensor-parallel path's bit for bit: once the planned rows are gathered, every rank runs the same chunk loop over the same rows against the same vocab-sharded head (test/runtime/distributed/test_qcp_prompt_logprobs.py asserts torch.equal against that path). Head TP over the query shards (--attn-head-tp-size equal to the shard group) adds the head exchanges — all-to-all transposes, permutations of bytes — and the o_proj tail. What can differ from the TP8 prefill baseline is therefore the output projection alone: its form, and for the row-parallel form the order in which the per-rank head partials are summed.

Head TP with the row-parallel o_proj computes the same per-rank partials TP8 does, but reduce-scatters them to the shard rows where TP8 all-reduces. Under rl-bitwise (--batch-invariant-collectives) both collectives are the ordered fold (ranks 0..W-1 left to right in fp32, one rounding; comm_backend/auto.py: route), so QCP with head TP sums o_proj exactly as the TP8 engine does, on any GPU set, and within the QCP engine the sharded extend (reduce-scatter) and the drafter's replicated decode steps (all-reduce) agree likewise. Known gap: the GPU validation of this layout against a plain TP4 prefill engine matched tokens and logprobs bitwise on prompts within index_topk, while one prompt whose context exceeded it — the indexer's top-k selecting a strict subset of the history — kept the tokens but diverged in logprobs from position 0 (max |Δ| 3.96e-2); unresolved, so the statement above is validated within index_topk only.

The head-replicated default (one GEMM over every head, the TP1 / trainer form) and head TP with --tp-batch-invariant attn (full-K column-parallel GEMM, a transpose back) have no cross-rank sum in o_proj at all and reproduce the TP1 / decode-side batch-invariant form — the one the RL trainer alignment wants, and the same gap to TP8 as between the decode side's batch-invariant layout and a cross-rank sum. The drafter's decode steps on a sharded engine merge partials across the KVP page owners (the page-sharded KV of --kv-parallel-size; combine_attention_partials, every head under the head-replicated layout, the attention-TP slice under head TP), with the ordered fold under rl-bitwise.

Kernel selection ​

The registry's two matching mechanisms split the work:

  • Traits are seller-declared: a kernel declaring deterministic={True}, batch_invariant={True} documents itself, and a requested trait excludes only kernels that declare the opposite. Good for ranking, useless for guarantees.
  • Features are buyer-required (subset test, silent kernels excluded): a leaf that affirmatively declares features={"batch_invariant"} is the only kind a batch-invariant request can be served by.

Batch-invariant leaves register the batch_invariant feature at the ordinary kernel boundary. The numerical envelope requires capabilities, not a provider name. Dense/LM-head and router projections, MLA absorption, sparse attention and expert execution bind their eligible implementations during initialization. MoE leaves additionally declare the requested routing and combine order; --moe-backend auto leaves provider selection to that constrained plan, while an explicit backend restricts its candidates without bypassing the capability check. Weight preparation and execution share the selected implementation. Missing capabilities fail planning instead of silently falling back. A capability declaration does not promise identical bits between different providers or agreement with a trainer; those comparisons need numerical acceptance for the chosen implementation and operation order.

Initialization-time kernel plans ​

A numerical requirement is independent of its provider. Dense projections and absorbed MLA request the batch_invariant feature through a KernelPlan. Initialization filters by formats, hardware and static model geometry, then ranks once, with the kernel name breaking equal-score ties. No representative batch is used. Candidates constraining a variable dimension are excluded: its exact values or the registry's _min, _max, and _align bounds. Static semantic traits such as topk_layout remain independent of a variable topk dimension. Planning a bounded deployment domain is not currently supported.

Execution validates the same formats and static traits and calls the saved implementation. It does not consult the registry, retune, or use an unplanned GEMV/fused fallback. A missing provider or an incompatible input fails. The provider's declaration covers its internal choices across the varying rows. plan_dsa_topk and plan_mla_project_value offer the same contract for the DSA indexer top-k and the MLA value projection: the index-key format, head geometry and slot order are static, while request counts and history lengths vary without provider selection. These operators run inside the model, so the model builds their plans and passes them in (kernel_plan= on dsa_prefill_topk / dsa_decode_topk, plan= on mla_project_value). No in-tree model plans them yet, since none is verified under rl-bitwise (see the per-model verification below); out-of-tree models that are verified are the current callers. Plans for different operators or phases still require cross-phase numerical acceptance; an individual capability declaration does not establish that.

A planned attention leaf that can be recorded into a CUDA graph also declares the cuda_graph feature. MLAAttnBackend and DSABackend derive their graph support from the leaves they planned (CudaGraphSupport.from_plans), so a leaf without the feature runs eagerly instead of being captured; startup logs which backend disabled which graph.

Absorbed MLA ​

Under a bitwise envelope, MLAAttnBackend plans the standard mla_extend_with_kvcache and mla_decode_with_kvcache operators with the batch_invariant requirement. Initial prefill and cached extend both prepare absorbed paged metadata, independent of the prefix length or query-length heuristic used in auto mode. The model supplies absorbed queries and writes the latent cache through its prepared-input path.

The current absorbed extend interface cannot express a sliding window, so this path rejects windowed configurations rather than widening their mask. Value projection stays separate to preserve its explicit rounding boundary. Provider registration owns dtype, head geometry and hardware constraints. The ordinary auto envelope retains its existing dynamic selection policy.

Explicit projection plans ​

A model may bind an existing unquantized Linear or LM-head method to a dense GEMM plan without changing its weight owner or storage. The logits processor uses that method for a registered bitwise envelope, preserving its normal input dtype cast and vocabulary communication.

Plugins may call register_numerics_envelope from the runtime registration surface. A policy supplies its name, bitwise classification, launch resolver and checkpoint validator. Plugin discovery precedes CLI choices and launch validation; registrations update shared containers and roll back together if plugin loading fails. Every policy still passes the model profile gate. Only rl-bitwise inherits the RL switches and unquantized/vocab-block limits; a plugin owns its own policy without changing individual execution entry points.

The existing gemm/mm plan also accepts already quantized FP8 operands and canonical scales. Quantization stays in the caller's existing quantization hook; an online quantizer is not selected by a frozen GEMM plan. The ordinary FP8 linear method retains its original dynamic selection and has no new plan field or forwarding path.

The original RMSNorm module may accept a preselected callable through kernel. The caller uses the existing registry selection during initialization; forward invokes that callable with the ordinary weight, epsilon, residual and output arguments. No plan or forward-scoped override is required. The default remains the original platform facade, and empty inputs keep the original early return. A custom callable rejects fused collective/norm entry points so its arithmetic cannot be bypassed. Kernel-specific input validation stays in the selected leaf.

MoELayer retains the usual moe/apply selection and its existing scoped kernel_override mechanism; no model-specific solution or feature arguments are added. An optional input dtype describes the actual activation for that selection without changing the global default dtype. A live-row mask is forwarded only to a kernel declaring that capability; unsupported masks fail rather than being ignored. Weight allocation and checkpoint loading stay on the original MoE layer for both precomputed and in-kernel routing.

Logprobs: one arithmetic for prompt and output ​

A returned logprob is the launch's --logprob-order applied to the row -- gather_token_logprobs in sampling/utils.py, the one function both consumers call: log_softmax(logits, -1, dtype=float32) gathered at the token under torch, fixed-width blockwise reduction under blockwise (alignment.trainer above) -- with the logits produced by the same LM-head route the sampler takes (LogitsProcessor._get_logits: quantized or dense GEMM, the planned batch-invariant GEMM under rl-bitwise, the TP gather, softcap). The sampler's output logprobs and the prompt (input) logprobs of the SGLang dialect (LogitsProcessor.compute_input_token_logprobs, requested through return_logprob + logprob_start_len) share that function, so the logprob of one token is the same number whether it was scored as a prompt position or sampled as an output -- the property an RL trainer relies on when it rescores a rollout. The dtype=float32 form widens bf16 logits inside the kernel (an exact conversion) instead of materializing an fp32 copy of the [rows, vocab] tensor first; because both paths go through the one function, whatever rounding the kernel applies is applied to both.

This comparison requires identical logits. With an ordinary unplanned GEMM, changing the prompt chunk size can change projection rounding; selecting a logprob order alone does not make that GEMM batch invariant. Dispatch tests use exactly representable projections to isolate the logprob operation, while rl-bitwise serving relies on its planned batch-invariant LM head.

Vocabulary-parallel argmax packs the shard maximum and its global index into an int64 key, reduced with MAX over the vocabulary group. This portable Triton kernel runs on CUDA, ROCm and Ascend; signed zeros have one canonical key so ties select the lowest global index. Ascend prepares this state before graph capture and uses HCCL for the reduction. Full-vocabulary callers retain the ordinary gather, dtype and sampling arithmetic above.

Accepted stop tokens and length limits are applied on the device before recurrent-state commit and accepted-frontier updates. Ascend uses the same Triton stop-condition kernel during eager and NPUGraph execution; the scalar reference is kept off captured device streams.

Top-k alternatives (top_logprobs_num, torch order only) follow the same rule: top_token_logprobs, beside it, takes a row's top_k largest entries of that fp32 log-softmax for the prompt rows and the prefill-sampled rows alike, so an alternative's value is the token's gathered logprob and, at the same top_k, a position's alternatives do not depend on which consumer scored it.

The TP gather differs in one respect that is not numerics: the sampled rows may take the multicast all-gather, whose result is a view of the group's shared buffer (safe because a whole forward separates consecutive calls), while the prompt-row chunks ask _get_logits for a private full-vocab tensor (require_full_vocab=True) and gather through the NCCL collective -- a chunk's log-softmax may still be reading its result when a faster rank issues the next chunk's gather. The gathered values are identical either way.

Prompt logprobs are gathered in position chunks of --input-logprob-chunk-tokens rows so the transient [rows, vocab] logits stay bounded. The chunk size is a sizing knob, not a numerics one: the reductions involved are row-local (the GEMM's row is independent of its neighbors under the per-row GEMM leaves of invariance.batch, and log-softmax reduces within a row), so no value depends on which chunk, or how large a chunk, a position landed in. The same holds across prefill chunks: positions are scored by the chunk that feeds them and assembled per request afterwards, so chunked prefill and prefix-cache hits do not change a prompt logprob either -- the admission probe is capped at logprob_start_len (scheduler.md §1) so every scored position is actually recomputed.

The current trainer defaults to negative cross-entropy, -(log(sum_exp) - target), yielding -0.0 for a certain token; rl-bitwise folds torch to blockwise to match these bits and preserves either explicit blockwise choice. This corrects the previous assumption that the trainer defaults to positive zero. blockwise-direct uses target - log(sum_exp) for trainers that require +0.0. Both use the same reduction; torch is unchanged. The shared logprob dispatcher rejects unknown orders, including the old megatron name, so programmatic callers cannot silently switch arithmetic after the rename.

Prefix hits resume at any token: KV bitwise, recurrent state within tolerance ​

A prefix hit resumes a prompt at the token count a previous request stopped at -- its prompt end, its finish, or a chunk end -- not at a page boundary (docs/design/cache-concepts.md, "Token-level prefix hits"). For attention layers this changes nothing: a query row depends only on the K/V rows it attends to, and chunked prefill already starts forwards at arbitrary positions, so a KV-only model produces the same tokens and logprobs with a hit at any T as it does cold, bitwise, and the invariance harness asserts that.

Recurrent layers (GDN, KDA, Mamba2) are different: their chunked kernels group tokens in kernel chunks counted from each forward's first token, so resuming from a stored state at a T that is not a multiple of the kernel chunk regroups a floating-point reduction that the cold forward computed in one grouping. The stored state is the exact state of the forward that wrote it, not a reordering of the cold one, and the two agree to the kernel's intra-chunk bf16 rounding (measured on GDN: state max abs difference ~2e-2 on a 0.7 scale, mean ~5e-4). The engine does not cut chunks or publication points to the kernel chunk; hits at any T are served. The oracle for a hybrid model is therefore greedy token agreement over the generation plus allclose of the resumed state against the cold scan at the hit, not torch.equal; a hybrid run that needs the bitwise claim of this document across a prefix hit must resume at a kernel-chunk multiple -- which the rl-bitwise flows do by construction, their prompts being the complete prefix of a finished sequence. Within one forward the recurrent kernels keep their batch invariance: a request's state does not depend on its co-batch.

Retraction is a byte copy: the bitwise continuation oracle ​

A capacity retraction suspends a request with its image -- every cache group's data blocks, the partial last page included, and the slot-state blob: the next step's inputs and draft candidates, the verifier's recorded draft distributions q, the drafter's stash, the sampling backend's per-request state (docs/design/cache-concepts.md, "Retraction image") -- and the restore copies those bytes back into fresh Device pages and a fresh request-pool slot. Nothing is recomputed and no forward runs in between for that request; its continuation reads only the restored bytes, its tokens and its (seed, position)-keyed sampler. Under rl-bitwise, where the kernels are batch-invariant, a forced retraction followed by a restore therefore yields exactly the tokens and output logprobs of an unretracted run, whatever batch the restored request lands in; without batch invariance the restored run differs only by the batch-dependent rounding any other batch composition would produce. The only request-keyed Device rows not in the image are token-derived (the committed-token history, the n-gram tail) and are reseeded from the control plane's token list, as on any slot handoff.

This is the test oracle for the suspend/resume path on both retracting roles: run a fixed prompt set under rl-bitwise twice, once plainly and once with --debug-force-retraction-interval N so every request is retracted at least once (several times for the longest), and assert output_ids and output_token_logprobs equal element-wise. The set includes a prompt-logprob request, an MTP configuration with --enable-speculative-sampling at temperature 1 (so the stash and the q path are exercised), a Prefilling victim through a negative interval, and on PD the decode node forced. An L3 prefetch that lands short is outside this oracle and needs none: it happens before the request is admitted, and the pages it did not land are computed by the ordinary prefill.

Expert placement and online rebalancing ​

An expert placement (--ep-num-redundant-experts, --init-expert-location) decides which rank computes which route. Under the rank-order MoE combine each rank's leaf returns a partial over its local slots and the host folds the partials in rank order, so the placement decides which routes land in which partial and a different placement moves the fold's rounding: the output is a function of the placement. A static placement is fixed per deployment and keeps the run-invariance contract; --enable-eplb makes the placement traffic-dependent state, so under the rank-order combine a rebalanced deployment is not run-invariant — not only across a rebalance, but across runs that rebalanced differently. The envelope therefore requires a placement-independent combine with --enable-eplb: every route computed on exactly one rank by a row-invariant leaf and the routes of a token folded in slot order (--moe-combine-order slot), which makes each route's value a pure function of the token and its logical expert's weights — replicas are byte-identical copies — and the output bitwise identical across rebalances. The envelope folds that combine in (alignment.trainer above), so the combination is accepted; --enable-eplb with the rank-order combine stays an auto launch. The load counters, the CPU-side algorithm and the P2P copies never enter the arithmetic; the dispatch algorithm stays static (required under rl-bitwise already) and drafts stay trivially placed.

Acceptance ​

An envelope is verified end to end, not per switch, and rl-bitwise has two harnesses, both run against a deployment launched with --numerics rl-bitwise. The invariance harness generates with returned logprobs for the same prompts (a) alone at bs=1, (b) packed with random co-batches, (c) across repeated runs, and asserts torch.equal on token ids and logprobs — base model and speculative decoding each, greedy and at temperature=1.0 with fixed seeds (the T>0 case is what sampling_stream=per-request exists for). A deployment that passes may advertise the contract; one that fails it has a bug, not a tolerance.

Trainer alignment is earned against a second program, so its harness has a reference the engine does not produce: a trainer dump. For a fixed prompt set (the invariance prompts plus longer, chat-formatted ones), the trainer runs its forward over each prompt + response sequence and records, per sequence:

ids          int64[T]       prompt followed by response token ids
logprob      float32[T-1]   log p(ids[t+1] | ids[:t+1]) at every position,
                            from the trainer's own log-softmax
layout       str            the trainer's TP/PP/EP factorization and dtype
commit       str            trainer commit that produced the dump

(an optional logits[:8, :] of the first positions helps bisect a mismatch to a layer). The engine then teacher-forces the same ids under --numerics rl-bitwise with return_logprob, logprob_start_len=0 and asserts torch.equal on the fp32 input_token_logprobs vector, reporting the first divergent position otherwise. Two self-consistency checks ride along: the response part of the teacher-forced input_token_logprobs must equal the output_token_logprobs the engine reported while generating that response, and the result must hold alone and packed exactly as for rl-bitwise. Agreement up to a tolerance is a finding to bisect, not a pass. The portable block_sumexp leaf and the vendor leaf reproducing the trainer's in-block order differ at ulp level, so the vendor leaf is what the comparison runs with.

The leaf contract behind --logprob-order blockwise is tokenspeed_kernel.ops.sampling.blockwise_logprobs's: a leaf registers ("sampling", "block_sumexp") with the signature format_signatures("shifted", "dense", {torch.float32}) and the batch_invariant feature (the op requires it), is called as leaf(shifted, *, block_size) with the fp32 [rows, vocab] logits already shifted by the row max, and returns fp32 [rows, vocab // block_size] block partials in vocabulary order, each block reduced in its own fixed tree. The op folds the columns left to right in fp32 — the trainer's rank-ordered reduction when the selected block width is its shard width. Pairing narrower shards first is a different reduction tree; do not infer cross-TP agreement from the leaf's pairwise in-block tree. The in-tree torch_block_sumexp sits at Priority.PORTABLE; the vendor leaf registers above it (any higher band), so plain selection takes it wherever it is installed and no host code names the vendor. The final subtraction follows the explicit form: blockwise negates the NLL and blockwise-direct subtracts log(sum_exp) from the target. Their signed-zero contract is described in the logprobs section above.

The pins above cover only the paths a verified model takes, so the verification is recorded per model and enforced at startup (require_verified_numerics): a model profile lists the envelopes its model passes in ModelProfile.numerics_envelopes, and launching an envelope other than auto refuses any target or draft model that does not list it — every in-tree model included, since none has a profile. The declaration is the model's promise and the harnesses are what keep it honest: a model declares rl-bitwise when it is ready to run them, and a declared model that fails either has a bug to fix, not a flag to set. Two incompatibilities are refused regardless of the declaration: quantized checkpoints (no batch-invariant quantized GEMM leaf exists, so their linears would select shape-dependent ones) and a vocabulary that is not a multiple of BLOCKWISE_VOCAB_BLOCK (the envelope's logprob order folds fixed blocks of it; MoE TP other than 1 is refused by resolve_numerics for the same reason on the MoE side).

Shared torch-order correction-bias routing ​

torch_router_topk in the runtime re-exports the kernel package's torch_softmax_bias_topk. Its arithmetic is unchanged: FP32 softmax, sorted PyTorch top-k on probabilities plus correction bias, gathering unbiased probabilities, scaling, then mapping identity candidates to -1. Exact ties retain PyTorch's order, not an independently imposed expert-ID order.

Callers needing a frozen choice use plan_softmax_bias_topk and execute softmax_bias_topk with that plan. The registered batch-invariant domain is NVIDIA SM90+ with at most 2048 candidates; larger widths fail planning rather than inheriting an unvalidated claim. The direct torch entry keeps the existing runtime behavior on other platforms and widths. Tests preserve the old formula independently and compare weights/IDs byte-for-byte on ties, cutoff ties, bias, identity candidates, near ties, batch packing/reordering and CUDA graph replay. This does not promise cross-version PyTorch tie compatibility.

Input and output vocabularies ​

The runtime uses ModelConfig.input_vocab_size for input token bounds and ModelConfig.vocab_size for the output vocabulary. The checkpoint convention is unchanged: its vocab_size is the embedding size and its optional output_vocab_size is the head size. A missing or null output_vocab_size means the two sizes are equal. Both must be positive integers, and the output vocabulary must fit within the input vocabulary for subsequent decoding. Invalid sizes are reported with both runtime field names and their values. The original Hugging Face config is not rewritten.

Input buffer clamps, n-gram token validation and multimodal draft substitute validation use input_vocab_size. Sampling, draft probability tensors, grammar masks, logit bias and prompt-logprob validation keep using vocab_size. Prompt-logprob admission checks only the scored next-token targets, input_ids[logprob_start_len + 1..], against the output vocabulary. Unscored context retains the input bounds, including the token at logprob_start_len; it feeds the first scored row rather than being that row's target. An input-only token can reach the model without becoming a valid output token or prompt-logprob target. Models or plugins with a reduced head must construct the embedding and head for their respective checkpoint sizes.

After loading a target model, initialization checks the logits processor's logical vocabulary against the runtime output size. A mismatch fails before sampling buffers or CUDA graphs are built; TP shards and padded weight rows are not logical vocabulary sizes. Workers without logits and draft models (which may use their own vocabulary mapping) are outside this target check.

Speculative input preparation bounds only draft candidate columns by the output vocabulary; the prompt and the first (verified) decode token retain input bounds. The existing defensive clamp still prevents a stale proposal from indexing outside output probability tensors. This extra clamp only runs for split vocabularies, stays on the device and is shared by eager and graph execution.

TokenSpeed documentation