diff --git a/x/mlxrunner/speculate.go b/x/mlxrunner/speculate.go index 3e66d5d8f..f5d278d56 100644 --- a/x/mlxrunner/speculate.go +++ b/x/mlxrunner/speculate.go @@ -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 } diff --git a/x/mlxrunner/speculate_stats.go b/x/mlxrunner/speculate_stats.go index 7425be876..68ad82342 100644 --- a/x/mlxrunner/speculate_stats.go +++ b/x/mlxrunner/speculate_stats.go @@ -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