CHAPTER 10

Different correctness contracts

The previous chapter separated two questions: whether mask and compact compute the same truncated model, and whether that model is close enough to a full prefill. A passing test answers the question encoded in its assertion. Changing the question can change the verdict without changing the engine.

python3 labs/21_equality.py

The experiment makes the reduction order depend on the number of rows. It then keeps that order fixed, and finally tries two other checks:

row-count-dependent: single=1, batched=0
row-count-dependent: masked=0, compact=1
fixed order: A=1 in every layout; 6 subset/order checks
adjacent float32 inputs: equal=False, abs diff=1.16415321827e-10
exp(d)-1-d metric=0; zero does not establish equality
logprob lists differ; both perplexities=7.389056099

The unrelated row changes the first result from 1 to 0. Fixing the order removes that difference. Yet a numerical difference metric can still return zero for unequal inputs, and two different logprob lists can have equal perplexity. These are three different observations, not three ways to prove the same claim.

Scope and order. We first explain the reduction, then follow DeepInfer’s batch-invariance flag into its libraries. Golden tokens let us compare computations along a fixed history. A second experiment separates repeatable random draws from a correct output distribution. Finally, we trace how grammar masks and matcher state constrain speculative output. The released tests supply the questions and assertions throughout; this chapter inspects them but does not run their model workloads.

10.1 Equality needs a fixed computation order

In exact arithmetic both ways of grouping the row give 1. In Python’s usual binary64 arithmetic, adding 1 to $10^{16}$ loses the 1. The lab’s single-row path adds from left to right:

Its multi-row path gives alternating entries to two lanes. The first lane adds $10^{16}+1$, the second contains $-10^{16}$, and the lane totals add to 0. A reduction combines several values into one; its grouping is part of the numerical computation.

Here is the entire layout-dependent rule. Masking keeps the batch shape but skips unwanted outputs; compacting removes those rows before the call:

19def reduce_rows(rows, active, fixed=False):
20    lanes = 1 if fixed or len(rows) == 1 else 2
21    result = {}
22    for name, values in rows:
23        if name in active:
24            partials = [left_sum(values[lane::lanes]) for lane in range(lanes)]
25            result[name] = left_sum(partials)
26    return result
27

Neither operation changes the retained row’s values. Nevertheless, two physical rows select two lanes even if only one output is wanted. One compact row selects one lane. With fixed=True, both calls use one lane and agree. The lab checks all six combinations of nonempty retained subsets and row orders for its two-row input, comparing masked and compact calls against a direct sum.

This isolates one cause of a batch-dependent result. An engine also has to align dtypes, fused operations, kernel choices and the visible attention keys. Removing query rows can preserve the remaining computation when rows are independent in those respects. Removing keys from a retained query changes its mathematical input, as chapter 9 showed. No choice of addition order can undo that change. Figure 10.1 sets this comparison beside the other two execution changes this chapter tests.

We will distinguish four kinds of claim. An exact comparison requires equal observed values. A tolerance check admits a specified difference under a specified metric. A distributional check asks whether sampled outputs follow the intended probabilities. A task or format check asks whether an answer satisfies a requirement, such as a grammar. A valid format does not establish a factually correct answer.

Several released comments call direct comparisons “bitwise”. We will state the comparator: Python == and torch.equal compare numerical values, not raw byte strings. For example, numerical equality does not distinguish positive and negative zero. More importantly for this chapter, a zero threshold on a computed metric need not even establish value equality. We return to the lab’s counterexample when we open test_kl.py.

10.2 One flag reaches several libraries

DeepInfer’s README offers a batch-invariance setting: a request’s result should not depend on its batch neighbours. The worker translates that setting into library configuration before running model computations [1, 76]. At lines 67–69 of deepinfer/worker/runtime_config.py, the condition is “batch invariance enabled, or NPU”. Its body is:

68deep_gemm.use_deterministic_algorithms(True)
69tile_kernels.use_deterministic_algorithms(True)

The NPU condition does not mean the complete feature is available there. Lines 61–63 separately assert that DEEPINFER_BATCH_INVARIANT is off on NPU. The adjacent comment says to remove this restriction after full support; the comment on line 67 says that the NPU setting prevents unstable CI numerics (author’s translations). These are distinct scopes: deterministic algorithms in particular libraries, and an engine-wide batch-invariance promise. The README describes the latter as NVIDIA-only.

The flag also reaches attention. The two sparse MLA decode calls in deepinfer/model_executor/models/v4_modules/attn.py, lines 666 and 1002, pass batch_invariant=envs.DEEPINFER_BATCH_INVARIANT. Earlier in the same file, it changes the choice of fused normalisation and rotary-position operations. Read the entire condition:

87self.use_fused_norm_rope_attn_rope_cast = (
88    # disabled in dense packed-KV decoding
89    ((not envs.DEEPINFER_BATCH_INVARIANT and cfg.engine_type == EngineType.PROMPT) or self.is_sparse or not self.has_packed_kv)
90    and self.num_heads in (64, 128)
91    and PLATFORM.is_gpu
92)

When the flag is on, EngineType.PROMPT alone no longer enables this fusion. But is_sparse or not has_packed_kv can still do so, subject to the head-count and GPU conditions. Saying that the flag disables all fused prefill paths would miss those alternatives. Aligning prefill and decode numerics is a plausible reason for controlling fusion; the predicate alone does not prove that explanation or the resulting equality.

Chapter 7 introduced a second source of layout changes: expert replicas. DeepInfer’s DeepEP buffer has this option:

79allow_multiple_reduction=not (envs.DEEPINFER_BATCH_INVARIANT and envs.DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL > 0),

The restriction requires both batch invariance and a positive expert rearrangement interval. Recall that replica selection includes a token’s index in the batch. Changing neighbours can therefore change its destination rank. A changed route may change how partial results are summed. That is a control-flow inference linking the replica rule to this option; these callers do not reveal or prove the libraries’ internal reduction order.

This chapter’s comparison is a division of responsibility, not a claim that DeepInfer invented deterministic reductions. DeepGEMM, TileKernels, FlashMLA and DeepEP are independent DeepSeek repositories. Chapter 7 records the pinned evidence for the replica rule and the caller’s communication paths. The attention calls lead to FlashMLA, named in chapter 3: DeepInfer’s attention backend imports its sparse attention and fused norm/RoPE operations, while the V4 attention module also imports the fused operations directly [27, 22]. The manifest records FlashMLA commit 2e5429fc, pinned by DeepInfer’s vendor/FlashMLA submodule. This identifies the dependency; it does not audit its kernel internals. A flag at any one boundary cannot establish the whole engine contract; the next section examines what the released tests actually compare.

10.3 Golden tokens hold the trajectory fixed

Suppose two runs disagree at token 7. By token 8, their histories differ, so their model scores need not agree even in exact arithmetic. To locate the numerical difference, a test can force the same token history while continuing to compute the model’s probabilities. This is teacher forcing. DeepInfer calls the forced sequence golden_token_ids.

The batch-invariance test first generates a batch of 64 requests. It then replays each request alone with the batched output as its golden sequence:

99for token_ids, output in zip(prompt_token_ids, outputs, strict=True):
100    sp = copy.copy(sampling_params)
101    sp.golden_token_ids = [x.token_id for x in output]
102    output_single = llm.generate(None, sp, [token_ids], use_tqdm=False)[0]
103    result_batched = _get_result(output)
104    result_single = _get_result(output_single)
105    assert result_batched == result_single, f"{dp_rank=}, {result_batched=}, {result_single=}"

The helper _get_result collects token IDs, their logprobs, and the top-5 IDs and logprobs. The assertion compares the whole tuple. However, equality of the output token IDs is imposed by the override. The substantial numerical check is whether the model assigns the same logprobs and top-5 results along that history. It is not a comparison of two independent stochastic draws [77].

10.3.1 What the override changes

The verification path in Sampler._verification first computes candidates, then calls get_accept_result. Its reference helper makes the override visible:

237accept = raw_accept
238tokens = candidate_token_ids
239if golden is not None:
240    use_golden = golden != -1
241    accept = torch.where(use_golden, golden == draft_token_ids, accept)
242    tokens = torch.where(use_golden, golden, tokens)
243blocked = ~(is_verify & accept)
244prefix = torch.cat([blocked.new_zeros(1, dtype=torch.int32), blocked.cumsum(0, dtype=torch.int32)])
245alive = prefix[:-1] == prefix[ftp[row_to_req.long()].long()]
246effective_accept = alive & is_verify & accept

A golden value of $-1$ means no override. Otherwise, the candidate becomes the golden token and acceptance becomes the test “draft equals golden”. The prefix calculation then prevents a request from accepting drafts beyond its first blocked row. The CUDA implementation repeats the same override before scanning preceding rows in deepinfer/model_executor/layers/tilelang_ops/cuda/sampler.py, lines 252–276 [78].

For a golden continuation $[a,c]$ and drafts $[a,b]$, the first draft matches, the second does not, and the output still follows $[a,c]$. The model’s stochastic acceptance probabilities did not decide that trajectory. This is why a golden-token test is useful for numerical comparison and unsuitable as evidence that rejection sampling preserves a distribution. Also keep these testing tokens separate from DFlash’s leading “golden” input: chapter 8 used that name for a committed token anchoring proposal generation.

10.3.2 The released inventory

The following checks vary different things. Test filenames beginning with test_ are under tests/models/; the references identify each fixed-version file [77, 80, 82, 75].

The vision test deserves its own row: it compares image embeddings before language-model sampling. It checks one image against repeated and mixed-shape batches using torch.equal. Passing a language-output comparison would not by itself identify whether that earlier representation was unchanged.

10.3.3 Decode versus prefill: inspect the metric

test_kl.py uses MTP 5, context parallelism 2 and expert parallelism 8, with image input and batch invariance enabled. Its text tokens and golden continuation are drawn uniformly from non-special tokens, not from the model distribution. Sampling uses temperature 0.5 and top-p 0.95. The first pass decodes that continuation. The second appends it to the prompt and obtains its logprobs through prefill. This holds the history fixed while changing how it is computed.

The compared values are processed logprobs. Temperature changes the scores, top-p keeps a nucleus of high-probability tokens, and the remaining tokens receive zero probability. The sampling reference shows the order:

209logits.mul_((1.0 / transform_inputs.temperatures).unsqueeze(1))
210apply_top_p_top_k_ref(logits, transform_inputs.top_ps, transform_inputs.top_ks)
211logprobs = torch.log_softmax(logits, dim=-1)

The top-p helper replaces excluded logits with $-\infty$, so their logprobs stay $-\infty$. Prompt logprobs use the same request settings: csrc/engine/prepare_sample.cpp, lines 186–188, copies temperature, top-p and top-k into the prompt arrays. The prompt-logprob helper passes those arrays to fused_sample before looking up each prompt token [77, 59].

The original prompt’s logprob lists are compared directly. The continuation’s logprobs are converted to float32 tensors and passed to cal_diff; the test accepts its fourth field at a threshold of zero. Before either comparison, it replaces $-\infty$ with $-9999$. Those transformations are part of what the test checks. If a token lies outside both nuclei, both values become the same sentinel, irrespective of their excluded logits. Such a position records agreement on exclusion, not agreement between finite logprobs. If nucleus membership differs between the runs, that normally produces a sentinel-versus-finite discrepancy.

Uniformly chosen tokens could make these exclusions common when the model’s nucleus contains few vocabulary entries. The fraction depends on the model and histories; it has not been measured here, and the test does not report it. To assess numerical coverage, record positions excluded on both sides, positions retained on both sides, and membership disagreements before replacing infinities. Finite-position differences and exclusion agreement answer different questions.

Two comparisons help locate this limitation. YOCO L2 uses temperature 0.5 and top-p 1.0, leaves top-k disabled and asserts that DEEPINFER_SAMPLER_FORCE_TOPK is zero: there is no nucleus or top-k truncation in that setup. The batch-invariance test instead obtains its golden tokens by sampling the batched run, rather than uniformly choosing vocabulary entries. Its chosen tokens therefore come from that run’s retained distribution; the single-request replay must match their logprobs and top-5 results. Neither contrast makes the selected probabilities an independent oracle for the whole model [75, 77].

The continuation comparison has a second limitation, in the metric itself:

4def cal_diff(x: torch.Tensor, y: torch.Tensor) -> tuple[float, float, float, float]:
5    if torch.equal(x, y):
6        return 0.0, 0.0, 0.0, 0.0
7    x, y = x.double(), y.double()
8    rmse = ((x - y) * (x - y)).mean().sqrt().item()
9    cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)
10    chebyshev_dis = (x - y).abs().max().item()
11    kl = ((torch.exp(x - y) - 1) - (x - y)).mean().item()
12    return cos_diff, rmse, chebyshev_dis, kl

Write $d=x-y$. The last field averages $e^d-1-d$ over the selected tokens. A KL divergence between two categorical distributions would require summing over their token probabilities with the appropriate weights. This expression is a discrepancy measure on the supplied logprobs; its name does not make it that distributional divergence.

In real arithmetic, $e^d-1-d$ is non-negative and zero only at $d=0$. Floating-point subtraction changes the conclusion. Lab 21 uses adjacent float32 values, both exactly representable in Python’s binary64 floats:

The lab evaluates the same scalar expression and gets zero. The original PyTorch helper also returns zero for these inputs on CPU despite torch.equal being false. The first return branch is sufficient for zero, but it is not necessary. To assert exact value equality, compare the values directly; to allow a tolerance, choose and state a metric that measures the intended difference.

This comparison completes part of chapter 5’s recomputation story. Rebuilding decoded history through prefill should reconstruct compatible representations and probabilities. The test probes that boundary for its selected inputs and settings. It does not actually preempt and resume a request, and its zero metric does not prove identical continuation. Sampling state, auxiliary state and the cache must also be compatible. Nor does this test enable chapter 9’s truncated-prefill mode.

10.3.4 Two other baselines

YOCO L2 compares mask with compact under batch invariance, fixing the golden trajectory and checking the other relevant settings. It directly compares chosen logprobs, top-k IDs and logprobs, and optional DFlash confidence values. The context-parallel derivation comment in deepinfer/model_executor/input_metadata/yoco.py, lines 136–139, explicitly points to an end-to-end equality check. Neither a local derivation nor agreement between these two truncated modes measures their error against full prefill. That remains chapter 9’s separate quality question.

The offline-inference test uses a different baseline: stored decode perplexities in tests/models/offline_decode_ppl.json. It checks the current run against those scalars, rather than comparing two fresh runs. The baseline is tied to the platform, workload and parallel configuration: NVIDIA uses EP 4 / CP 2; Ascend uses EP 16 / CP 8. Both have two data-parallel ranks. Changing that configuration requires a new baseline.

For chosen-token logprobs ${\ell }_1,\dots ,{\ell }_n$, perplexity is $\text{exp}(-{\text{∑}}_i{\ell }_i∕n)$. Lab 21 compares $[-1,-3]$ with $[-2,-2]$. Both give $e^2$, although neither token’s logprob agrees. Exact equality of an aggregate is an aggregate regression check, not tokenwise equality.

10.4 Repeatable and correctly distributed are different claims

Numerically equal probabilities are only part of repeatable sampling. A draw also consumes randomness. A shared stream hands its next value to whichever request happens to run next. Moving a request changes the value it receives.

python3 labs/22_sampling_contracts.py

shared stream AB: A=0, B=1
shared stream BA: A=1, B=0
keyed draws: all 6 batch orders and single requests agree
q=1, position 7: fresh target, token=1
q=3, position 7: accepted draft, token=0
ideal marginal in both budgets: [1/2, 1/2]

The shared stream changes both requests’ tokens when their order is reversed. The keyed version survives all six orders of three requests and each request run alone. But changing a verification budget still changes one token, even though the toy target and draft have identical distributions.

10.4.1 Keys belong to requests, positions and operations

The teaching generator hashes an explicit request seed and a position/role key. It uses SHA-256, not Python’s process-randomised hash, and returns an exact rational in $[0,1)$:

19def keyed_draw(seed, position, role, draft=False):
20    effective_seed = seed + (DRAFT_OFFSET if draft else 0)
21    message = f'{effective_seed}:{3 * position + role}'.encode('ascii')
22    bits = int.from_bytes(sha256(message).digest()[:8], 'big')
23    return F(bits, 2**64)

DeepInfer uses its device RNG instead. The following line numbers refer to deepinfer/model_executor/layers/tilelang_ops/cuda/fused_sample.py. In that file, the effective seed is the request’s seq_seed plus a namespace offset. The wrapper sets that offset to zero for the default path and to $1,000,000,007$ for MTP. The CUDA kernel uses three position-derived sequences [79]:

Here $p$ is the kernel’s sample_positions field. The seed is selected at line 227. Batch slot is not part of these RNG keys.

There is an important intermediate transformation. For packed verification rows, process_sampling expands the per-request parameters using row_to_req and local_step. Its reference expansion is:

61req = row_to_req.long()
62return SamplingKernelInputs(
63    LogitsTransformInputs(params.logits.temperatures[req], params.logits.top_ps[req], params.logits.top_ks[req]),
64    params.seq_seeds.view(torch.int64)[req].view(torch.uint64),
65    params.sample_positions[req] + local_step * DRAFT_POS_HASH_MUL,
66)

DRAFT_POS_HASH_MUL is $10,000,019$, not 1. Thus $p$ is not always a literal absolute token index. Reordering requests while preserving their local work preserves this construction; changing local verification positions need not. DFlash separately forms proposal positions from the base, the last valid offset and the draft index, then samples in the MTP seed domain. The lab deliberately uses simple absolute positions to isolate the role/domain change. It reproduces neither this position mixing nor the device generator.

10.4.2 Why the budget matters even with a fixed seed

Chapter 8’s timing curve influences the total verification budget and each request’s query count. The sampling metadata determines the role of each row:

18row_to_req = torch.repeat_interleave(
19    torch.arange(ftp.shape[0], dtype=ftp.dtype, device=ftp.device),
20    seqlen_q,
21    output_size=total_q,
22)
23local_step = torch.arange(total_q, dtype=ftp.dtype, device=ftp.device) - ftp[row_to_req.long()]
24is_verify = local_step != seqlen_q[row_to_req.long()] - 1
25return row_to_req, local_step, is_verify

For a query count $q=1$, local row 0 is the final row and takes a fresh target sample. For $q=3$, rows 0 and 1 verify drafts, while row 2 draws afresh. At the same output position, the first run may use a target draw and the second an accepted proposal drawn earlier in the MTP namespace. Equal seeds do not make those two operations consume the same randomness.

Lab 22 sets target and draft probabilities to $(1∕2,1∕2)$, making every draft acceptable. Its budget loop follows that final-row rule:

39def budget_output(seed, start, q):
40    # q-1 verified drafts, then one fresh target draw, as in the metadata rule.
41    output, roles = {}, {}
42    for local in range(q):
43        position = start + local
44        verify = local != q - 1
45        if verify:
46            draft = token(keyed_draw(seed, position, 0, draft=True))
47            assert keyed_draw(seed, position, 2) < 1  # target/draft ratio = 1
48            output[position], roles[position] = draft, 'accepted draft'
49        else:
50            output[position] = token(keyed_draw(seed, position, 0))
51            roles[position] = 'fresh target'
52    return output, roles

With seed 11 at position 7, a fresh target draw produces token 1; the accepted draft produces token 0. Under ideal uniform draws each rule still has exactly the distribution $(1∕2,1∕2)$. Equality of marginal distributions does not require equality for a particular seed. The SHA-256 example supplies a concrete repeatability counterexample; the marginal statement follows from the two sampling rules, not a claim that a finite list of hash outputs proves uniformity.

This gives a source-based explanation for the README recipe: enable DEEPINFER_BATCH_INVARIANT, disable DEEPINFER_ENABLE_MTP_VARLEN, and provide a fixed SamplingParams.seq_seed on NVIDIA. Measured timings can change the budget; the budget can change the sampling operation and its keys. This does not mean every timing perturbation changes a token.

No end-to-end test of that complete recipe was found in the inspected tests/models/ inventory: none sets seq_seed. Embedded sampler and RNG kernel tests are different checks. The evaluation driver supplies per-sample seeds but initially sets varlen to 1; its later –env loop can override that value. Its default setup is therefore not the README’s recipe. Seeds alone, or the helper’s repeatability docstring, do not establish that those evaluation runs are repeatable [79].

10.4.3 Test the probabilities, not just the seed

test_sampler.py asks whether the sampled tokens agree statistically with their reported probabilities. It uses MTP 5 and four settings $(T,\text{top-p})=(0,1),(1,0.95),(8,0.95),(1,1)$, with 128 trajectories per setting and up to 2,048 output tokens per trajectory [80]. It does not force a golden continuation.

A categorical sample is not a uniform number. A randomised probability integral transform (PIT) turns it into one under the target distribution. Order the tokens, give token $x$ an interval of width $p_x$, and let $F(x)$ be the interval’s upper end. Given the sampled token, draw an independent uniform $U$ and form

If token $x$ is chosen with probability $p_x$, that probability is spread uniformly over an interval of the same width. The intervals together fill $[0,1]$ uniformly. The construction works conditionally on each token’s history; the probabilities can change between positions.

DeepInfer calls the upper endpoint sum_smaller_probs. It sorts by ascending probability, breaking ties by column index, and includes the sampled token. Its reference calculation makes the tie rule explicit:

38pivot = logprobs.gather(1, sample_results.unsqueeze(1))
39mask_less = logprobs < pivot
40col_idx = torch.arange(logprobs.shape[1], device=logprobs.device).unsqueeze(0)
41mask_equal_before = (logprobs == pivot) & (col_idx <= sample_results.unsqueeze(1))
42out = (logprobs.exp() * (mask_less | mask_equal_before)).sum(dim=1)
43out.masked_fill_(invalid, 0.0)
44return out

Lab 22 uses target probabilities $(1∕2,1∕4,1∕4)$. The token order is 1, 2, 0, with intervals $[0,1∕4]$, $[1∕4,1∕2]$, and $[1∕2,1]$. A biased sampler instead chooses token IDs 0, 1 and 2 with probabilities $(3∕4,1∕8,1∕8)$. Exact interval calculations give:

PIT quarter masses, correct: 1/4, 1/4, 1/4, 1/4
PIT quarter masses, biased:  1/8, 1/8, 3/8, 3/8

The correct sampler places a quarter of its PIT mass in each quarter of the unit interval. The biased sampler places too much in the upper half. These are exact rational mass calculations, not a histogram or a run of DeepInfer’s statistical test.

The released test filters the arrays before constructing PIT values and calling a Kolmogorov–Smirnov (KS) check:

83ssp_arr = np.array(ssp, dtype=np.float64)
84p_arr = np.exp(np.array(log_probs, dtype=np.float64))
85if not (np.all(np.isfinite(ssp_arr)) and np.all(np.isfinite(p_arr)) and np.all(p_arr > 0)):
86    return None
87if (ssp_arr - p_arr < -0.01).sum() > 0:
88    return None
89rng = np.random.RandomState(seed)
90u = np.clip(ssp_arr - p_arr * rng.uniform(0, 1, size=n), 0, 1)
91ks_stat, ks_quantile = ks_test_uniform(u.tolist())
92if ks_stat is None:
93    return None
94return ks_stat, ks_quantile, n

KS measures the largest gap between the sample’s cumulative distribution and the uniform CDF. The helper turns that gap into an approximate p-value using an asymptotic formula with a finite-sample correction. A p-value describes how unusual such a discrepancy would be under the null model; it is not the probability that the sampler is correct. The helper returns a p-value under the name ks_quantile. A second KS check compares those trajectory-level values with a uniform distribution.

The actual failure gate is short:

242if len(all_quantiles) == 0:
243    errors.append(f"DP {dp_rank} {config_name}: no valid trajectories for all, skip={skip_counts}")
244elif all_pv is None:
245    errors.append(f"DP {dp_rank} {config_name}: computing the all p_value failed")
246elif all_pv <= 1e-5:
247    errors.append(f"DP {dp_rank} {config_name}: all p_value={all_pv} is too small; the sampler may have a problem")

all_pv uses the whole trajectories. The test also prints next1_pv, based on tokens after steps marked by draft_tok, but does not reject on that value. The filtering is consequential: if any position has $\mathtt{\text{ssp}}-p<-0.01$, the helper returns None and that whole trajectory is omitted, not failed. An inclusive CDF should satisfy $\mathtt{\text{ssp}}\ge p$; the tolerance filter can therefore discard evidence of the very probability inconsistency being investigated. The helper also omits missing or non-finite data and non-positive probabilities, then clips the surviving PIT values to $[0,1]$. If fewer than two valid trajectory-level values remain, the aggregate is unavailable and the test fails. Otherwise, the aggregate may pass despite discarded trajectories.

Passing means this finite test did not detect a discrepancy at its threshold. It cannot prove exact distributional equality or exclude chapter 8’s small, deliberate residual penalty. There is another uncovered claim: PIT compares samples with the distribution the engine reports after its transformations. If temperature or top-p is applied wrongly, both sampling and the reported probabilities could use the same wrong transformed logits and still pass. An independent check of the transformations would be needed to detect that error. This test also does not enable structured output, which adds another transformation of the distribution.

10.5 A grammar narrows the distribution

A grammar allows only certain continuations of the text already accepted. DeepInfer represents the allowed token IDs as a bitmask. Its sampler applies that mask to the logits before sampling and draft verification:

428ftp = decode_inputs.mtp_first_token_positions
429seqlen_q = decode_inputs.mtp_seqlens_q
430row_to_req, local_step, is_verify = prepare_verification_metadata(ftp, seqlen_q, total_q, self.max_MTP_value)
431speculative_inputs = driver_inputs.speculative
432
433# Apply structured-output masks before sampling and draft verification.
434if so_indices is not None:
435    for step, indices in enumerate(so_indices):
436        if len(indices):
437            apply_token_bitmask_by_step_inplace(logits, so_bitmasks[step], indices, ftp, step)

For an allowed set $S$, simple masking and renormalisation would turn a distribution $p$ into $p(x)∕{\sum }_{y\in S}p(y)$ for $x\in S$, and zero elsewhere. That describes a probability contract only after defining the order of masking, temperature, top-k/top-p and speculative correction. A format test alone cannot establish this distribution. DeepInfer places the grammar mask before process_sampling, whose transformations and verification we traced in chapter 8.

10.5.1 The mask belongs to a tentative history

The grammar compiler and matcher are provided by the external XGrammar library, pinned as xgrammar==0.2.2. DeepInfer’s wrapper in deepinfer/structure_output/matcher.py owns trigger progress, bound-token state, EOS bookkeeping and rollback. Unlike the independent DeepSeek kernel repositories, XGrammar is an external project supplying grammar machinery. DeepInfer supplies the integration with speculative requests [81].

On the host, Sampler._get_so_mask first advances the matcher with the previous step’s accepted tokens and final output. Status 1 marks initial setup; status 2 enables subsequent updates:

351try:
352    if so_status[idx] == 2:
353        for token in seq_metadata.last_accepted_token + [seq_metadata.last_output_token]:
354            if not matcher.accept_token(token):
355                raise RuntimeError(f"matcher rejected step token {token}, upstream invariant broken")
356    elif so_status[idx] == 1:
357        so_status[idx] = 2
358    else:
359        raise RuntimeError(f"unexpected so_status {so_status[idx]}")

Now suppose a teaching grammar permits ab!, ad!, cb! and cd!, with one character per token. Before any output, the mask allows $\{a,c\}$. Proposed drafts $[a,b]$ require two later masks: after $a$, allow $\{b,d\}$; after $\mathit{ab}$, allow $\text{{!}}$. Those masks assume a prefix that has not yet passed target verification (Figure 10.2).

DeepInfer fills the mask for position 0, then tentatively advances through the drafts. This is the start of that loop:

366num_accepted_steps = 0
367for k, draft_token in enumerate(seq_metadata.last_next_n_token):
368    if matcher.accept_token(draft_token):
369        if matcher.constrained():
370            mask_filled = matcher.fill_vocab_mask(self.bitmasks_cpu[k + 1], idx)

The omitted lines 371–372 assert that filling succeeded and register the request’s mask index. The loop then counts successful tentative accepts and restores the matcher:

373        num_accepted_steps += 1
374    else:
375        break
376matcher.rollback(num_accepted_steps)

If verification rejects $a$ and selects $c$, the next host update commits $c$, not $\mathit{ab}$. If it accepts $a$ but rejects $b$, it commits $a$ and the replacement token. The saved mask for a later row is usable only along the speculative prefix for which it was built; accepted-prefix logic prevents later candidates from becoming output after an earlier rejection.

The wrapper’s rollback capacity is nextn+2. Its rollback method restores the inner matcher’s state and the wrapper’s trigger/EOS/bound bookkeeping. The extra two positions are an implementation allowance; the constant alone does not establish why precisely two are required. What matters for this path is that tentative advancement and restoration stay paired.

10.5.2 Failure ends the request

A compilation future can time out or fail. Its wrapper returns an aborted matcher. When the engine enqueues a request with such a matcher, it marks FINISHED_SO_FAILED in deepinfer/engine/llm_engine.py, lines 621–622. During sampling, a missing or aborted matcher sets so_status to $-1$; runtime exceptions also mark it aborted and set that status. The scheduler closes the path by checking the stored status:

976if (seq.structured_output.status == -1) {
977    seq.schedule.status = SequenceStatus::FINISHED_SO_FAILED;
978    seq.append_token_id(0);
979    seq.append_compressed_token_id(0);
980    token_cnt++;

The two appended zero IDs are the scheduler’s failure bookkeeping. The branch marks the sequence failed instead of continuing unrestricted generation. The preceding comment explicitly says to consult the stored status, not only the report for this step: a step with no active structured-output work may report no status (author’s translation). A missing report must not erase an earlier failure.

This is also the boundary of chapter 5’s preemption discussion. The scheduler sets so_blocks_recompute whenever structured-output status is nonzero. Such a request terminates on preemption even if ordinary requests would be recomputed. Restoring KV would not, by itself, restore the matcher and its trigger history. Reconstructing that state would be a requirement for a new recovery implementation, not a feature present in this release. The speculative rollback above is implemented; preemption recovery of the matcher is not.

10.5.3 Where the constraint begins, and what the tests cover

A trigger can leave an initial part of the output unconstrained. Once its token sequence matches, the grammar takes over. A bound token has a more specific meaning: if it fires while the trigger is pending, the wrapper forces the remaining trigger tokens and then enters the grammar. It is not an ending delimiter that releases the constraint. These rules live in accept_token and fill_vocab_mask, with their state included in rollback.

The released structured-output example supplies three prompts asking for incompatible formats, then requests a regex, a JSON schema and an EBNF grammar. All use temperature zero. It requires nonempty output, a normal stop finish reason and a final EOS token before decoding the text.

The format checks are explicit:

54if schema_type == "regex":
55    assert re.fullmatch(so_args["schema"], text), f"regex: {text!r}"
56elif schema_type == "json":
57    value = json.loads(text)
58    assert isinstance(value, dict) and set(value) == {"color", "count"}, f"json: {text!r}"
59    assert value["color"] in {"red", "green", "blue"}, f"json: {text!r}"
60    assert type(value["count"]) is int and value["count"] in {1, 2, 3}, f"json: {text!r}"
61else:
62    assert schema_type == "ebnf"
63    assert re.fullmatch(r"answer=(yes|no);level=[1-3]", text), f"ebnf: {text!r}"

The JSON case checks keys, allowed values and a strict integer type. The EBNF case checks the resulting text with a full-match regex. The offline-inference test explicitly invokes these examples. They establish these syntax and termination checks, not stochastic masked-distribution correctness, factual accuracy, or coverage of every trigger, timeout and rollback path. The sampling KS test does not use structured output, so its result cannot fill that gap.

10.6 Experiments to try next

1. Change lab 21 to keep two reduction lanes for every batch size. Does batch equality return? Does the result have to equal the one-lane result? Explain why numerical consistency and an exact-arithmetic answer differ.
2. Replace the zero-metric check with an assertion that matches an exact value-equality contract. Then propose a separate check for a stated absolute tolerance. What should either check do with non-finite values?
3. Construct two different three-token logprob lists with equal perplexity. Which test in the inventory could distinguish them directly?
4. For one request with four candidate rows, mark the verified rows and fresh row. Reduce its budget to two. Identify which operation changes at the second row and which parts of the RNG key are relevant.
5. In the grammar diagram, accept $a$, reject $b$, and choose $d$. Trace the matcher’s tentative state, rollback state and next committed state. What additional state would a preemption/recompute implementation need?
6. Complete this template for a released test and for a proposed end-to-end deterministic-sampling test. Include the model and hardware configuration; do not treat a recipe as an executed test.

You can now follow a DeepInfer correctness claim from its configuration through the relevant library or state machine to the assertion that checks it, distinguishing fixed history from sampled output, stable random keys from a stable sampling path, and grammar validity from the probability or quality of the answer. Those boundaries tell you which test to run when changing the engine—and when a new test is needed.