Skip to content

Design: batch-invariant greedy decode (opt-in fixed decode width) #80

Description

@iamwavecut

Problem

Greedy decoding is not bitwise-reproducible across batch compositions. Every token-batched decode op selects its kernel — and often its internal schedule — from the step's token count t = Σ_lanes (1 + accepted drafts). Variants differ in reduction order (split-K partition count, warp-level K partitioning, tile width), so logits differ by sub-ulp amounts between compositions; on near-ties the argmax flips and long greedy continuations diverge. Measured on RTX 5090 with tools/smoke/kv_cache_concurrency.py: a burst of 8 identical 256-token greedy requests produced 2–3 distinct continuations with the host cache off, 3 with it on; solo double-runs are identical. This is the documented contract today (docs/serving.md, Determinism).

Where t selects arithmetic (verified in the tree)

op selector t classes
ops::linear W8G32 select_w8_a16_launch(n,k,t) + launch_w8_small_t exact-t instantiations ladders per (n,k): ≤4 / ≤16 / ≤33–48 / mma; inside small_t: kTileTokens buckets 8/16/24/32/40/48 and kKWarps 8→4 at t>24 or >32
ops::linear Q4/Q5/Q6/BF16/FP8/NVFP4 per-family ladders incl. exact-t simt (q6 c5/c6/c7) and split-K exact (q5) same shape
attn_input_proj, linear_swiglu, linear_pair, linear_add (W8) route tables by problem.cols family-specific
GDN input projection (W8) DecodeR8Direct (1) / SplitKMmaDirect (2–96) / MmaR64C128 3 classes; NVFP4 adds DecodeFused/SmallTFused
GQA attention decode gqa_attention_uses_small_t(tokens ≤ 6) per-lane width per-lane (acceptance-dependent, not neighbour-dependent)

RoPE, norms, sampling are per-element/per-lane and invariant. Prefill is lane-exclusive and invariant. MoE (35B) is batch-dependent by construction and out of scope.

Measured (RTX 5090, sm_120, tests/ops/linear/test_t_invariance.cpp, column 0 fixed, t = 1…64)

New reduction-order classes begin at (bitwise change of out[:,0] versus the previous t):

shape class boundaries
W8 [1024×5120] t=17 (368 of 1024 elements, max
W8 [6144×5120] t=17 (2494 of 6144 elements, max
W8 [14336×5120] t=5 (3 of 14336 elements, max
W8 [34816×5120] t=24 (6 of 34816 elements, max
W8 [5120×6144] t=32 (1 of 5120 elements, max
W8 [5120×17408] t=32 (2 of 5120 elements, max
W8 [5120×10240] t=25 (2 of 5120 elements, max
W8 [248320×5120] t=33 (40 of 248320 elements, max
Q6 [248320×5120] t=8 (104076 of 248320 elements, max
Q4 [1024×5120] none — invariant for t ≤ 64
Q4 [6144×5120] t=17 (3 of 6144 elements, max
Q4 [34816×5120] t=2 (3 of 34816 elements, max

Reading: the boundaries coincide with the dispatcher ladders (4/16/24/32/48) and the small_t kKWarps 8→4 buckets; the differing elements are 1–2 bf16 ulps apart — exactly the near-tie flips that drive greedy drift. Note the early boundaries inside ordinary decode ranges: W8 [14336×5120] at t=5, Q6 vocabulary head at t=8, so even --max-concurrency 2 with MTP crosses them.

Cost caveat (supersedes the "≤2–3%" estimate)

With a fixed T_pad the whole model runs the T_pad-class kernels even at t=1; the padded MMA work is not free once T_pad reaches the compute-bound regime (27B at T_pad=32 is roughly balanced between weight bandwidth and MMA throughput on a 5090), so single-lane decode may slow by a large fraction under the flag. The flag therefore targets multi-lane servers whose real t already sits near T_pad, and must ship with measured t=1 / t=T_pad throughput.

Proposal: opt-in fixed decode width (--decode-invariant)

Default path untouched. In invariant mode every GEMM-class op on the decode path runs at a fixed token width T_pad = max_concurrency × (1 + draft_tokens) rounded to the family's tile: activations are padded (garbage columns are harmless — columns never mix), the kernel and its internal schedule are therefore identical for every real t ≤ T_pad, and the padded output columns are ignored. Invariance follows by construction for all column-independent ops; per-lane ops (attention, GDN recurrence) stay at their real widths and are already neighbour-invariant. Cost: extra MMA work on padding columns in a weight-bandwidth-bound regime (expected small; to be measured at t=1 and t=T_pad), plus fewer CUDA-graph profiles. Scope for a first PR: the 27B W8 decode path (linear + the four fused W8 projections + GDN W8 input). Acceptance: phase E of the concurrency battery asserts bitwise parity across the burst under the flag; solo runs unchanged. Not a default: pinning one schedule changes results of steps that previously crossed ladder thresholds, so bit-equality with existing outputs is not promised even in invariant mode.

Alternatives considered

  • Rounding/clamping logits: changes default outputs — rejected.
  • Batch=1 decode: destroys batching — debugging knob only.
  • Pinning launcher pointers (no padding): insufficient — small_t and fused plans re-dispatch internally on t.

Verification plan

tests/ops/linear/test_t_invariance.cpp (landed on the #73 branch) maps the classes per registered shape; the battery's phase E under the flag is the end-to-end assertion; determinism_fingerprint.py separates real nondeterminism from the composition effect.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions