Optimize iLLaDA end-of-thinking logit boost

[mercurystraw]

Summary

This PR optimizes the </think> logit-boosting logic used by variable-length iLLaDA generation.

The existing implementation scans generated tokens one position at a time with torch.equal. On CUDA, every comparison returns a Python boolean and therefore introduces a device synchronization. This becomes increasingly expensive for
long reasoning outputs.

The new implementation performs sequence detection and per-row token selection with batched tensor operations while preserving the existing boost behavior.

Existing behavior

During each diffusion decoding step, apply_end_think_logit_boost encourages the model to emit the token sequence corresponding to </think>.

The original implementation performs the following operations for each row:

1. Clone the complete [batch, sequence, vocabulary] logits tensor.
2. Scan every possible position in the generated sequence.
3. Call torch.equal for each candidate window to search for </think>.
4. Locate the first masked candidate position.
5. Inspect the preceding tokens to determine which token in </think> should
   be boosted next.
6. Add the progress-dependent boost to the selected logit.

The positional scan is implemented as a Python loop:

for start in range(sequence_length): torch.equal(...)

When the tensors are on CUDA, each torch.equal result must be synchronized back to Python. Since this check runs at every diffusion step, the number of host-device synchronizations grows quickly for long generations.

With steps=32 and block_length=32, a generation that does not emit </think> can perform approximately:

• 2.1 million positional comparisons at gen_length=2048
• 8.4 million positional comparisons at gen_length=4096

The original implementation also clones the full logits tensor before checking whether any row still needs a boost. For iLLaDA's vocabulary size, a batch-1 BF16 logits clone is approximately 1.2 GiB at a sequence length of 4096, excluding the additional prompt length.

Changes

Vectorized end-sequence detection

The positional Python loop is replaced with a batched tensor operation:

tokens.unfold(1, sequence_length, 1) == sequence

The resulting windows are reduced with all and any to determine whether each row already contains the complete </think> sequence.

This keeps the comparisons on the GPU instead of synchronizing once per candidate position.

Batched candidate and prefix handling

The implementation now computes the following values for all rows together:

• whether a masked candidate position exists;
• the first candidate position;
• whether the complete end sequence has already been emitted;
• which token in the end sequence should be emitted next.

For example, if </think> tokenizes into three tokens, the logic remains:

• no matching suffix: boost the first token;
• first token already emitted: boost the second token;
• first two tokens already emitted: boost the third token;
• complete sequence already emitted: do not apply a boost.

The longest matching prefix continues to take precedence, matching the original implementation.

Preserve the generation-context boundary

Complete-sequence detection is restricted to:

tokens[:, context_start:]

Partial-prefix matching also checks that the prefix does not cross context_start.

This prevents a sequence present in the prompt, or one spanning the prompt/completion boundary, from being treated as a generated </think>.

Avoid unnecessary logits cloning

Rows are marked active only when they:

• still contain a masked candidate position; and
• have not already emitted the complete end sequence.

If no row is active, the function returns the original logits immediately. The full logits tensor is cloned only when a boost actually needs to be applied.

Preserve numerical behavior

Boost progress is computed in float64 before being converted to the logits dtype. This matches the previous Python floating-point calculation and avoids small rounding differences that could affect low-confidence token ordering.

Only active rows are updated through indexed tensor assignment. The input logits remain unmodified whenever a boost is required, preserving the previous non-in-place behavior.

Expected impact

This change is expected to:

• remove the per-position Python/CUDA synchronization loop;
• reduce latency for long iLLaDA reasoning generations;
• avoid full-logits cloning after </think> has already been generated;
• reduce temporary allocation and CUDA allocator pressure;
• preserve generated-token selection and boost semantics.

The model forward pass, softmax, remasking strategy, generation schedule, end-think boost strength, and EOS handling are unchanged.

[Copilot]

Pull request overview

This PR refactors and optimizes the </think> logit-boosting path used during (variable-length) iLLaDA generation to avoid per-position Python loops and reduce CUDA host-device synchronization overhead.

Changes:

• Replaces per-row positional scanning (torch.equal in a Python loop) with batched sequence detection via unfold + reductions.
• Vectorizes “next token in end-sequence” selection across the batch while preserving context boundary handling (context_start) and longest-prefix precedence.
• Avoids cloning the full logits tensor when no rows are eligible for boosting.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[mercurystraw]

The change improved throughput from 11.1 to 17.8 tokens/s (~1.6×) and reduced total runtime from 26h to 15.5h on one H100 GPU in a Math500 local comparison, but it was not fully controlled. For reference only.