[Perf] Reuse DSpark probability distance across loss and metrics

[Levius-Fubuki]

[Perf] Reuse DSpark probability distance across loss and metrics

Purpose

DSpark currently computes the same draft/target probability distance once for
acceptance-rate metrics and again when the L1 distillation loss is enabled:

before: {softmax: 4, abs/sub vocabulary reduction: 2}
after:  {softmax: 2, abs/sub vocabulary reduction: 1}

This PR computes the per-token L1 probability distance once and reuses it for:

• accept_rate = clamp(1 - 0.5 * l1_distance, 0, 1);
• the weighted L1 distillation loss.

It removes one pair of vocabulary-sized FP32 softmax results and one duplicate
probability-difference reduction. The public loss API, loss definitions,
confidence targets, metric definitions, and distributed normalization remain
unchanged. The implementation stays on the portable PyTorch path with no
custom kernel, architecture gate, or device-specific tuning.

Test Plan

Regression and equivalence tests

python -m pytest tests/modeling/dspark/test_loss.py -vv

The tests verify:

• the L1-enabled path calls torch.softmax twice instead of four times;
• shared probability-distance values match the original expression;
• acceptance-rate values match the original expression;
• draft-logit gradients match the original expression;
• weighted L1 numerators and denominators are unchanged;
• L1 enabled/disabled and confidence enabled/disabled paths;
• no probability work runs without aligned target logits.

Repository test suite

python -m pytest -q
python -m py_compile \
  deepspec/modeling/dspark/loss.py \
  tests/modeling/dspark/test_loss.py
git diff --check

RTX 4090 loss-path microbenchmark

The standalone harness measures _collect_local_terms plus backward with L1
loss enabled. Baseline and candidate run from separate worktrees at the same
upstream base commit. Each row uses 20 warmups and 100 measured iterations.

Hardware/software: NVIDIA RTX 4090 (SM89), CUDA 13.0, PyTorch 2.12.1+cu130,
Python 3.12.3. No GPU configuration changes.

Vocabulary	Token rows	Old median / IQR (ms)	New median / IQR (ms)	Median speedup	Old / New peak allocation
151,936	128	3.592 / 0.025	2.531 / 0.128	1.42x	671.3 / 523.0 MiB
151,936	512	17.167 / 0.069	11.593 / 0.024	1.48x	2.61 / 2.03 GiB

For the 512-row case, the complete local loss collection and backward path is
approximately 32.5% faster and uses approximately 22.2% less peak allocated
memory.

An isolated probability-distance sweep covered vocabulary sizes 32,768,
65,536, and 151,936 with 128, 512, and 2,048 token rows. The larger shapes
showed 1.51x-1.57x median speedups, and all tested shapes completed without OOM.

Test Result

6 passed

The focused suite, repository-wide pytest invocation, Python syntax compilation,
and git diff --check passed locally.

Limitations

These are isolated DSpark loss-path measurements from one RTX 4090, not
end-to-end training throughput. No performance claim is made for other GPU
architectures, CPU execution, distributed training, or configurations with L1
loss disabled. The smallest isolated shapes were sensitive to GPU measurement
order; the reported whole-loss results use DSpark's large-vocabulary regime.

The benchmark harness is intentionally not part of this patch. It compares the
exact original and candidate expressions and separate upstream/candidate
worktrees; the production patch contains only the loss refactor and focused
regression tests.

AI assistance

AI assistance was used during implementation, testing, benchmarking, and PR
drafting. The human contributor reviewed and takes responsibility for the
change and evidence.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.