mlxrunner: count every speculative round and log stats at info

The per-request stats are the main diagnostic for speculative
throughput, so log them at info; the controller line stays debug.
Recording chosen depths at the next beginRound dropped rounds with no
successor, so record at endRound and count resume as a depth-0 round.
This commit is contained in:
Jesse Gross 2026-07-29 14:43:42 -07:00
parent 4713800b08
commit 3ff2dcb649
2 changed files with 11 additions and 7 deletions

View file

@ -147,7 +147,6 @@ func (s *speculation) open(request Request, caches []cache.Cache) *speculationSe
func (s *speculationSession) beginRound() {
now := time.Now()
if !s.lastRoundStart.IsZero() && s.roundDrafts >= 0 {
s.stats.recordRound(s.roundDrafts)
if s.roundDrafts == s.prevDrafts {
s.spec.depth.cost.observe(s.roundDrafts, now.Sub(s.lastRoundStart))
}
@ -162,6 +161,7 @@ func (s *speculationSession) beginRound() {
// positions past it (a terminator, not a target rejection).
func (s *speculationSession) endRound(drafted, accepted, observed int) {
s.roundDrafts = drafted
s.stats.recordRound(drafted)
s.stats.iterations++
s.stats.drafted += drafted
s.stats.accepted += accepted
@ -277,7 +277,10 @@ func (st *speculativeDecoder) resume() []sampler.Result {
st.position = position
st.inner.close()
st.inner = nil
// No round spans this call, so the next beginRound attributes no cost.
// The drained emission counts as one more plain round, but no round spans
// this call, so the next beginRound attributes no cost.
st.s.stats.recordRound(0)
st.s.stats.iterations++
st.s.roundDrafts = -1
return next
}

View file

@ -18,9 +18,6 @@ type specStats struct {
}
func (s *specStats) recordRound(depth int) {
if !slog.Default().Enabled(context.TODO(), slog.LevelDebug) {
return
}
s.chosen = append(s.chosen, depth)
}
@ -49,7 +46,7 @@ func (s *specStats) depthOverTime() string {
}
func (s *speculationSession) logStats() {
if !s.enabled || !slog.Default().Enabled(context.TODO(), slog.LevelDebug) {
if !s.enabled {
return
}
acceptance := 0.0
@ -62,7 +59,11 @@ func (s *speculationSession) logStats() {
avgDraft = float64(s.stats.drafted) / float64(s.stats.iterations)
avgAccepted = float64(s.stats.accepted) / float64(s.stats.iterations)
}
slog.Debug("speculative decode stats", "iterations", s.stats.iterations, "drafted", s.stats.drafted, "accepted", s.stats.accepted, "acceptance", fmt.Sprintf("%.2f", acceptance), "avg_draft", fmt.Sprintf("%.2f", avgDraft), "max_draft", s.stats.maxDraft, "avg_accepted", fmt.Sprintf("%.2f", avgAccepted), "depth_over_time", s.stats.depthOverTime())
slog.Info("speculative decode stats", "iterations", s.stats.iterations, "drafted", s.stats.drafted, "accepted", s.stats.accepted, "acceptance", fmt.Sprintf("%.2f", acceptance), "avg_draft", fmt.Sprintf("%.2f", avgDraft), "max_draft", s.stats.maxDraft, "avg_accepted", fmt.Sprintf("%.2f", avgAccepted), "depth_over_time", s.stats.depthOverTime())
if !slog.Default().Enabled(context.TODO(), slog.LevelDebug) {
return
}
// Log learned acceptance over the trusted positions [1, frontier] and
// expected throughput over the searched window [0, frontier+1]; deeper