[orestis-z]

Purpose

Split out of the measurements in #1259 / #1262 / #1263, one change per PR.

Every training step built its metric scalars with torch.tensor(<python float>, device=cuda), which copies from pageable host memory and so cannot be issued asynchronously — the host blocks until the work already queued on the device drains. The same step ran torch.masked_select on the accuracy mask, whose output size depends on the mask's values, forcing another device-to-host sync before the count could be read. Building those scalars with torch.full/torch.zeros/torch.ones and counting masked positions with a reduction gives identical numbers without either sync.

full_total and cond_total stay separate tensors so a caller reducing the metrics dict in place cannot count one of them twice.

Measured

Single H100, eagle3, synthetic batches, full 151936-token vocab, total_seq_len 2048, 10 warmup + 50 measured steps, 3 reps, median of the per-run means, log_freq=1 (the default), scripts/benchmark.py. This PR against main:

Qwen3-0.6B	Qwen3-1.7B	Qwen3-8B	Qwen3-32B
1.111x	1.110x	1.030x	1.013x

Where the gain comes from, change by change

Same harness and settings, each change measured against its parent commit, so every cell is one PR's marginal contribution at that verifier size. Run-to-run spread within a cell is under 1% everywhere except main at 0.6B.

PR		Qwen3-0.6B	Qwen3-1.7B	Qwen3-8B	Qwen3-32B
#1294	metrics off host	1.111x	1.110x	1.030x	1.013x
#1295	reuse the argmax	1.076x	1.026x	1.013x	1.007x
#1296	single-pass fused KL	1.047x	1.017x	1.012x	1.003x
#1297	logs off the thread	1.029x	1.042x	1.015x	1.006x
#1298	one metric readback	0.980x	1.004x	0.997x	0.999x
#1299	H2D prefetch †	1.024x	1.013x	1.029x	1.021x
#1300	skip dropped positions	0.997x	0.988x	0.980x	0.988x
#1301	precast frozen heads	1.055x	1.096x	1.059x	1.036x
	cumulative	1.358x	1.329x	1.141x	1.077x

† #1299's row is not the prefetch. That commit also pins the synthetic batch in scripts/benchmark.py, and the real DataLoader already sets pin_memory=True, so it is a benchmark-fidelity fix rather than a training speedup. Applying only the pinning hunk on top of #1298 gives 1.017x / 1.017x / 1.030x / 1.022x — the whole row. Discounting it, the cumulative row is 1.335x / 1.307x / 1.108x / 1.054x.

Reading it:

• perf(train): keep metric construction off the host #1294 and perf(train): keep frozen linear weights in the autocast dtype #1301 carry most of the gain. They are the two largest at every size, and perf(train): keep frozen linear weights in the autocast dtype #1301 is the only one that also moves peak memory (−1.5 GB at 0.6B to −7.4 GB at 32B).
• perf(eagle3): reuse the argmax the forward already computes #1295, perf(losses): compute the fused KL and RKL forwards in a single pass #1296 and perf(train): render console logs off the training thread #1297 fade with verifier size but stay positive — 1–2% at 8B, 0.3–0.7% at 32B.
• perf(train): read the step metrics back with one device-to-host copy #1298 is a tie at every size. One sync per step instead of a few dozen is the right shape, but it does not show up here.
• perf(train): copy the next batch to the device on a side stream #1299's prefetch is inside noise in this regime once the harness fix is separated out, though it does pay once the metric log emit is deferred — see perf(train): copy the next batch to the device on a side stream #1299.
• perf(eagle3): skip the projection for positions the loss drops #1300 is negative at seq_len 2048 at every size and costs 0.3–1.5 GB of peak; it needs both long sequences and a small hidden size to pay.

Only #1300 and #1301 change peak allocated memory at all; the other six are flat to the megabyte.

Most of these changes remove fixed per-step host and metric overhead, so the relative win shrinks as the verifier grows even though the absolute saving does not. At 8B and 32B the step is dominated by the optimizer (42% and 58%, all Muon Newton–Schulz), which nothing here touches — the same point #1111's side note made.

Tests

pytest tests/unit/models/test_metrics.py

New cases cover that the metric tensors are built without a host-to-device copy and that the masked counts match masked_select. Passes.

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

• Source mode: Auto
• Repository: vllm-project/speculators
• Branch: perf/metrics-no-host-sync
• File: .coderabbit.yaml
• Commit: 2d9e141
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: e635c0f3-b1eb-4a79-8b77-26f04a861350

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: vllm-project/speculators
• Branch: perf/metrics-no-host-sync
• File: .coderabbit.yaml
• Commit: e845fed
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: 156b9557-4d43-48c5-a993-32dd83298ec9

📥 Commits

Reviewing files that changed from the base of the PR and between 94db8b8 and e845fed.

📒 Files selected for processing: 6

• src/speculators/losses/utils.py
• src/speculators/models/eagle3/core.py
• src/speculators/models/eagle3/metrics.py
• src/speculators/models/metrics.py
• src/speculators/train/trainer.py
• tests/unit/models/test_metrics.py

Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 3 remain after this review.

Walkthrough

The changes update scalar tensor initialization in loss, metric, and training code. They also change masked accuracy counting and add tests for masked counts, empty masks, and total tensor independence.

Changes: Loss and metric tensor handling

Layer / File(s)	Summary
Scalar tensor initialization src/speculators/losses/utils.py, src/speculators/models/eagle3/core.py, src/speculators/models/eagle3/metrics.py, src/speculators/train/trainer.py	Scalar loss, metric, and error-record tensors now use torch.zeros, torch.ones, or torch.full instead of torch.tensor constructors. The changes retain the specified device and dtype.
Masked accuracy counts src/speculators/models/metrics.py, tests/unit/models/test_metrics.py	compute_accuracy_single_step computes correctness and totals using the loss mask without selecting masked elements first. Tests cover masked counts with and without prev_correct, an all-false mask, and independent total tensors.

Possibly related issues

• vllm-project/speculators#1251: This PR changes metric accounting for masked positions, but does not add agentic rollout ingestion.

Priority: Normal

Merge Risk: Minimal · previous review e845f

The changed metric construction and accuracy counts have no established regression in the supported training path. The PR is mergeable after normal checks.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: keeping training metric construction off the host.
Docstring Coverage		Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 10 functions across 6 files.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description check		The description directly matches the changes. It explains the device-side metric construction, the replacement of masked selection with reductions, the separation of metric tensors, benchmark results,…

Full details: Description check

Explanation

The description directly matches the changes. It explains the device-side metric construction, the replacement of masked selection with reductions, the separation of metric tensors, benchmark results, and tests.

✨ Finishing Touches: 💡 1

📝 Generate docstrings

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 Require approval from approved reviewers list

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = fynnsu
  • approved-reviews-by = dsikka
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz

[fynnsu]

Looks good, but maybe cleanup comments.

Also as we discussed, I think there's an open question of whether we should also be applying the torch.full() switch to all uses of torch.tensor() with scalar values in the repo, or keep this pr scoped to Eagle3.