perf(eagle3): skip unused teacher projections while preserving metrics

[julyanghar]

EAGLE3 currently normalizes/projects every verifier hidden row before applying the loss mask. This draft projects only the rows required by the loss and the existing cross-depth accuracy history, then scatters into the unchanged dense targets layout. The draft backbone, draft logits/tokens, and metric definitions are unchanged.

Draft because the performance benefit is workload-dependent: contiguous prompt/response masks benefit, but dense masks and some short dispersed masks regress. This is not yet a recommendation to enable the change for every workload.

The important correctness constraint is that loss_mask alone is insufficient. prev_correct begins at supervised draft starts and compares start r against teacher row r+t at depth t, including masked intermediate targets. Keep the union of those positions for t in [0, ttt_steps). A three-position [1,0,1] regression test shows that zeroing masked targets alone preserves loss but changes later accuracy counts. The implementation preserves that existing behavior rather than changing metrics as part of a performance patch. With no mask or all rows needed, it retains dense projection (the latter still pays mask-analysis overhead).

Validation on base 3419401a380db305ed6493ba4680ad8a3c3e0e45:

• CUDA_VISIBLE_DEVICES='' python -m pytest tests/integration/models/test_eagle3_target_projection.py -q -o addopts='': 106 passed. Seven built-in eager losses, five mask layouts, depths 1/3/4, packed documents; real model loss, complete output metrics, draft tokens, and parameter gradients compared against dense teacher projection. Includes the mask-only negative control.
• CUDA_VISIBLE_DEVICES=2 python -m pytest tests/integration/models/test_eagle3_target_projection.py tests/integration/models/test_model_forward.py -k 'eagle3 or Eagle3' -q -o addopts='': 50 passed, 100 deselected, before expanding the new loss parameterization from KL to all seven losses. Includes existing compiled CUDA EAGLE3 forward/backward and attention-backend checks.
• CUDA_VISIBLE_DEVICES='' python -m pytest tests/unit/models/test_metrics.py tests/unit/models/test_eagle3_attention.py -q -o addopts='': 46 passed.
• Repository-pinned Ruff 0.15.20 check/format and git diff --check passed for changed files.

Benchmark scope: teacher norm/projection/selection/scatter only, not complete training. RTX 6000 Ada; PyTorch 2.11.0+cu130; BF16; B=1, H=1024, V=32768, TTT=7. Each arm wrapped with torch.compile, three warmup calls, four alternating AB/BA pairs, five calls per arm, synchronized host wall timing without profiler. Compilation is outside timing. Dense zero-filled targets remain allocated. All required-row logits were bit-identical in these sampled eager BF16 checks and argmax agreed; this is not a universal bitwise-equivalence claim.

Rows	Layout	Supervised rows	Retained teacher rows	Baseline ms	Candidate ms	Time change / baseline
2048	suffix	204	204	0.687	0.412	-39.99%
2048	dispersed	204	1025	0.803	0.832	+3.60%
2048	all	2048	2048	0.685	0.740	+8.02%
2048	empty	0	0	0.691	0.237	-65.73%
8192	suffix	819	819	2.792	1.049	-62.41%
8192	dispersed	819	4183	3.320	3.177	-4.31%
8192	all	8192	8192	3.701	3.983	+7.62%
8192	empty	0	0	2.748	0.701	-74.50%

The mask's future-position union explains why ~10% dispersed supervision can require ~50% of teacher rows. Nonzero/dynamic selection adds synchronization and compiled graph boundaries, so the full-mask fallback is not free. Eager function measurements also showed a short dispersed-mask slowdown (~13%). Full training throughput, distributed training, and serving/quality metrics have not been measured.

Related: #1088 replaces the metrics contract; if it lands first, the required-row rule and its tests need to be revisited. This PR targets current main and does not duplicate the metric-definition changes in #1088, #624, or #902.

Reproduce the function-only benchmark from the PR checkout

Save as a Python file and run with the desired CUDA_VISIBLE_DEVICES. It writes timing JSON, environment/base identity, and a working-tree patch into teacher-projection-benchmark/. The PR diff supplies the committed patch.

import json
import platform
import subprocess
import time
from pathlib import Path
import torch
from speculators.models.eagle3.core import Eagle3DraftModel

root=Path('teacher-projection-benchmark')
root.mkdir(exist_ok=True)
repo=Path.cwd()
class Head(torch.nn.Module):
    _compute_training_targets=Eagle3DraftModel._compute_training_targets
    def __init__(self):
        super().__init__()
        self.verifier_norm=torch.nn.RMSNorm(1024)
        self.verifier_lm_head=torch.nn.Linear(1024,32768,bias=False)
    def dense(self,h,m,t):
        return self.verifier_lm_head(self.verifier_norm(h))

provenance={'base':subprocess.check_output(['git','rev-parse','HEAD'],cwd=repo,text=True).strip(),'torch':torch.__version__,'gpu':torch.cuda.get_device_name(),'python':platform.python_version(),'scope':'torch.compile teacher norm/projection/gather/scatter function only; no complete training throughput','warmup':3,'pairs':4,'iterations_per_arm':5,'dtype':'bfloat16','H':1024,'V':32768,'depth':7}
(root/'provenance.json').write_text(json.dumps(provenance,indent=2))
(root/'speculators.patch').write_bytes(subprocess.check_output(['git','diff'],cwd=repo))
torch.manual_seed(42)
head=Head().cuda().bfloat16();results=[]
with torch.no_grad():
 for seq in [2048,8192]:
  h=torch.randn(1,seq,1024,device='cuda',dtype=torch.bfloat16)
  for layout in ['suffix','dispersed','all','empty']:
   mask=torch.zeros(1,seq,device='cuda',dtype=torch.bool)
   if layout=='suffix':mask[:,-int(seq*0.1):]=True
   if layout=='dispersed':mask[:,torch.randperm(seq,device='cuda')[:int(seq*0.1)]]=True
   if layout=='all':mask[:]=True
   needed=mask.clone()
   for t in range(1,7):needed[:,t:] |= mask[:,:-t]
   a=head.dense(h,mask,7);b=head._compute_training_targets(h,mask,7)
   av=a[needed];bv=b[needed];relative=float(torch.linalg.vector_norm(av.float()-bv.float())/torch.linalg.vector_norm(av.float()).clamp_min(1e-20))
   torch.testing.assert_close(av,bv,atol=0.02,rtol=0.02)
   assert torch.equal(av.argmax(-1),bv.argmax(-1))
   del a,b,av,bv
   torch.compiler.reset()
   funcs={'baseline':torch.compile(head.dense),'candidate':torch.compile(head._compute_training_targets)}
   for fn in funcs.values():
    for _ in range(3):fn(h,mask,7)
   measurements=[]
   for pair in range(4):
    for arm in (['baseline','candidate'] if pair%2==0 else ['candidate','baseline']):
     torch.cuda.synchronize();torch.cuda.reset_peak_memory_stats();start=time.perf_counter()
     for _ in range(5):out=funcs[arm](h,mask,7);del out
     torch.cuda.synchronize()
     measurements.append({'pair':pair,'arm':arm,'ms':(time.perf_counter()-start)*1000/5,'peak_allocated':torch.cuda.max_memory_allocated()})
   row={'seq':seq,'layout':layout,'loss_rows':int(mask.sum()),'required_rows':int(needed.sum()),'relative_L2':relative,'measurements':measurements};results.append(row)
   (root/'benchmark.json').write_text(json.dumps(results,indent=2))
   print({k:v for k,v in row.items() if k!='measurements'}, {a:sum(v['ms'] for v in measurements if v['arm']==a)/4 for a in funcs},flush=True)

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

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

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

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

This rule is failing.

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 = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz