perf(eagle3): gate the drafter quality diagnostics on the log level

tests/integration/test_eagle3_diagnostics_gating_contract.py

@@ -0,0 +1,152 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contract tests for the EAGLE-3 drafter quality diagnostics.
+
+They are log-only and every one of them forces a device sync, so they must not
+run on the hot path unless the log will actually be emitted.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import pytest
+
+QUALITY_LOG_PREFIX = "[drafter logits quality]"
+BACKEND_LOGGER = "verl_speco.backends.eagle3_trainer_backend"
+
+
+def _backend_and_inputs():
+    import torch
+    from omegaconf import OmegaConf
+
+    from verl_speco.backends.eagle3_trainer_backend import Eagle3TrainerBackend
+
+    vocab_size, seq_len, ttt_length = 6, 4, 2
+
+    backend = Eagle3TrainerBackend(
+        OmegaConf.create(
+            {
+                "rollout": {
+                    "drafter": {"training": {"use_logits": False, "ttt_length": ttt_length}}
+                },
+                "model": {"path": "/tmp/none"},
+            }
+        ),
+        OmegaConf.create({}),
+    )
+    backend.target_model = lambda last_hidden: torch.zeros(
+        *last_hidden.shape[:-1], vocab_size
+    )
+
+    class _FakeDraft:
+        t2d = torch.ones(vocab_size, dtype=torch.bool)
+
+        def __call__(self, **kwargs):
+            return {
+                "logits": [
+                    torch.randn(1, seq_len, vocab_size) for _ in range(ttt_length)
+                ],
+                "position_masks": [
+                    torch.ones(1, seq_len) for _ in range(ttt_length)
+                ],
+            }
+
+    batch = {
+        "input_ids": torch.zeros(1, seq_len, dtype=torch.long),
+        "hidden_states": torch.zeros(1, seq_len, 8),
+        "last_hidden_states": torch.zeros(1, seq_len, 8),
+        "attention_mask": torch.ones(1, seq_len, dtype=torch.long),
+        "loss_mask": torch.ones(1, seq_len),
+        "position_ids": torch.arange(seq_len).unsqueeze(0),
+    }
+    return backend, _FakeDraft(), batch
+
+
+class _Recorder(logging.Handler):
+    """Records everything the backend logger emits, whatever its level is.
+
+    ``caplog.at_level`` would raise the logger to DEBUG, which is exactly the
+    condition under test, so the handler is attached directly instead.
+    """
+
+    def __init__(self):
+        super().__init__(level=logging.DEBUG)
+        self.records: list[logging.LogRecord] = []
+
+    def emit(self, record):
+        self.records.append(record)
+
+
+def _compute_loss_at(level: int) -> list[logging.LogRecord]:
+    backend, model, batch = _backend_and_inputs()
+    backend_logger = logging.getLogger(BACKEND_LOGGER)
+    previous_level = backend_logger.level
+    recorder = _Recorder()
+    backend_logger.setLevel(level)
+    backend_logger.addHandler(recorder)
+    try:
+        backend.compute_loss(model, batch, 0)
+    finally:
+        backend_logger.removeHandler(recorder)
+        backend_logger.setLevel(previous_level)
+    return [r for r in recorder.records if QUALITY_LOG_PREFIX in r.getMessage()]
+
+
+def test_quality_diagnostics_are_silent_above_debug() -> None:
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+
+    assert _compute_loss_at(logging.INFO) == []
+
+
+def test_quality_diagnostics_are_emitted_at_debug() -> None:
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+
+    quality_records = _compute_loss_at(logging.DEBUG)
+
+    assert quality_records
+    # Routine per-step training metrics belong at DEBUG, not WARNING.
+    assert all(r.levelno == logging.DEBUG for r in quality_records)
+
+
+def test_quality_diagnostics_do_not_sync_above_debug() -> None:
+    """Above DEBUG the loss must not pay for the diagnostics' device syncs."""
+    torch = pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+
+    original_item = torch.Tensor.item
+
+    def run_at(level: int) -> int:
+        backend, model, batch = _backend_and_inputs()
+        backend_logger = logging.getLogger(BACKEND_LOGGER)
+        previous_level = backend_logger.level
+        backend_logger.setLevel(level)
+        calls = 0
+
+        def counting_item(self):
+            nonlocal calls
+            calls += 1
+            return original_item(self)
+
+        torch.Tensor.item = counting_item
+        try:
+            backend.compute_loss(model, batch, 0)
+        finally:
+            torch.Tensor.item = original_item
+            backend_logger.setLevel(previous_level)
+        return calls
+
+    assert run_at(logging.INFO) < run_at(logging.DEBUG)

verl_speco/backends/eagle3_trainer_backend.py

@@ -1178,6 +1178,10 @@ def compute_loss(self, model, batch, _current_pad_size):
             0.0, device=input_ids.device, dtype=torch.float32
         )
         gamma = 0.8
+        # The per-step quality stats below are log-only, and every one of them
+        # forces a device sync. Collect them only when the log will be emitted,
+        # matching how the EAGLE-1/2 backend gates the same diagnostics.
+        diagnostics_enabled = logger.isEnabledFor(logging.DEBUG)
 
         # Preprocess shifted targets
         for idx in range(length):
@@ -1234,7 +1238,7 @@ def compute_loss(self, model, batch, _current_pad_size):
                     position_mask=position_mask,
                 )
                 target_top1 = target_p.argmax(dim=-1)
-            if (
+            if diagnostics_enabled and (
                 base_valid_position.any()
                 and not valid_position[base_valid_position].all()
             ):
@@ -1244,7 +1248,7 @@ def compute_loss(self, model, batch, _current_pad_size):
                     int(dropped_tokens.detach().cpu().item()),
                 )
             with torch.no_grad():
-                if valid_position.any():
+                if diagnostics_enabled and valid_position.any():
                     draft_top1 = logits.argmax(dim=-1)
                     step_top1_correct = (
                         (draft_top1[valid_position] == target_top1[valid_position])
@@ -1303,7 +1307,11 @@ def compute_loss(self, model, batch, _current_pad_size):
             total_local_ploss += (gamma**idx) * step_loss_sum
             total_local_tokens += valid_position.float().sum()
 
-        if use_sparse_restricted_ce and sparse_base_tokens.detach().float().item() > 0:
+        if (
+            diagnostics_enabled
+            and use_sparse_restricted_ce
+            and sparse_base_tokens.detach().float().item() > 0
+        ):
             logger.debug(
                 "[drafter sparse restricted ce] base_tokens=%s valid_tokens=%s dropped=%s "
                 "intersection_mean=%.6f hit_mass_mean=%.6f min_intersection=%s min_hit_mass=%s",
@@ -1326,8 +1334,8 @@ def compute_loss(self, model, batch, _current_pad_size):
                 logits_sparse_min_mass,
             )
 
-        if quality_tokens.detach().float().item() > 0:
-            logger.warning(
+        if diagnostics_enabled and quality_tokens.detach().float().item() > 0:
+            logger.debug(
                 "[drafter logits quality] valid_tokens=%s top1_acc=%.6f top%s_acc=%.6f "
                 "local_ploss_sum=%.6f local_tokens=%s per_step=%s",
                 int(quality_tokens.detach().cpu().item()),