From 70af8532ae884bc9f69530562c5903b384c58db4 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 3 Nov 2026 23:20:43 +0000 Subject: [PATCH] feat(utils): add once-per-message logger methods debug_once / info_once / warning_once log on rank 0, once per distinct message and args. select_kernel, which logged its choice once per layer, and the hand-rolled log-once guards now use them. --- .../layers/quantization/linear/nvfp4.py | 15 ++++------- .../freetoken/layers/quantization/method.py | 3 ++- .../layers/quantization/moe/nvfp4.py | 4 +-- python/freetoken/models/gguf/mmproj.py | 7 +----- python/freetoken/models/minimax_m3/config.py | 25 +++++-------------- python/freetoken/models/qwen3_vl/vision.py | 14 +++-------- python/freetoken/utils/logger.py | 16 ++++++++++++ 7 files changed, 36 insertions(+), 48 deletions(-) diff --git a/python/freetoken/layers/quantization/linear/nvfp4.py b/python/freetoken/layers/quantization/linear/nvfp4.py index 1002a70b3c..7f222ce7b1 100644 --- a/python/freetoken/layers/quantization/linear/nvfp4.py +++ b/python/freetoken/layers/quantization/linear/nvfp4.py @@ -38,21 +38,16 @@ def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: return nvfp4_dense_linear_t(x, layer.weight, layer.weight_scale, layer.weight_global, layer.bias) -@functools.cache -def _warn_part_globals() -> None: - logger.warning_rank0( - "an NVFP4 linear fuses parts whose weight global scales differ by more than %d%%; the W4A4 kernels use the largest " - "for every part, which scales the other parts' outputs by their ratio; --quant-backend linear.nvfp4=triton keeps " - "each part's global", round((_PART_GLOBAL_WARN_RATIO - 1) * 100), - ) - - def _tensor_global(layer: Any) -> torch.Tensor: """The one weight global the W4A4 GEMMs and the small-batch kernels take: the largest per-part global, the block scales unchanged.""" g = layer.weight_global.float() g_max = g.max() if g_max > _PART_GLOBAL_WARN_RATIO * g.min(): - _warn_part_globals() + logger.warning_once( + "an NVFP4 linear fuses parts whose weight global scales differ by more than %d%%; the W4A4 kernels use the largest " + "for every part, which scales the other parts' outputs by their ratio; --quant-backend linear.nvfp4=triton keeps " + "each part's global", round((_PART_GLOBAL_WARN_RATIO - 1) * 100), + ) return g_max diff --git a/python/freetoken/layers/quantization/method.py b/python/freetoken/layers/quantization/method.py index 9e90637577..b8ea1afe49 100644 --- a/python/freetoken/layers/quantization/method.py +++ b/python/freetoken/layers/quantization/method.py @@ -38,7 +38,8 @@ def select_kernel(candidates: Sequence[type], requested: str, cfg: Any): continue if kernel.worth_it(cfg): if skipped: - logger.info("kernel %s selected; skipped %s", kernel.name, "; ".join(skipped)) + # every layer selects its own kernel + logger.info_once("kernel %s selected; skipped %s", kernel.name, "; ".join(skipped)) return kernel if fallback is None: fallback = kernel diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py index 2a5e8bacd0..96cd1884de 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -163,7 +163,7 @@ def _marlin_symbols_ok() -> bool: ) from vllm.scalar_type import scalar_types # noqa: F401 except Exception as exc: - logger.warning( + logger.warning_once( f"NVFP4 marlin backend is installed but unusable ({exc!r}); " "falling back to the Triton inline-dequant backend" ) @@ -371,7 +371,7 @@ def _b12x_symbols_ok() -> bool: prepare_w4a16_packed_weights, ) except Exception as exc: - logger.warning( + logger.warning_once( f"NVFP4 b12x backend is installed but unusable ({exc!r}); " "falling back to the Triton inline-dequant backend" ) diff --git a/python/freetoken/models/gguf/mmproj.py b/python/freetoken/models/gguf/mmproj.py index a85828164a..912711f2f7 100644 --- a/python/freetoken/models/gguf/mmproj.py +++ b/python/freetoken/models/gguf/mmproj.py @@ -126,15 +126,10 @@ def find_mmproj(model_path: str) -> Mmproj | None: return _find(model_path, os.environ.get(MMPROJ_ENV, "")) -@functools.cache -def _warn_inline_vision(model_path: str) -> None: - logger.warning("vision: %s carries vision tensors itself (Ollama layout), which are not read; use a llama.cpp mmproj file", model_path) - - def warn_inline_vision(model_path: str, tensor_names: set[str]) -> None: """Old Ollama GGUFs carry the tower inside the text file under another layout; FreeToken reads only llama.cpp mmproj files.""" if any(n.startswith(("v.", "mm.")) for n in tensor_names): - _warn_inline_vision(model_path) + logger.warning_once("vision: %s carries vision tensors itself (Ollama layout), which are not read; use a llama.cpp mmproj file", model_path) def mmproj_tensor(info: GgufTensorInfo) -> torch.Tensor: diff --git a/python/freetoken/models/minimax_m3/config.py b/python/freetoken/models/minimax_m3/config.py index bbecadd3b3..94d33cd2d5 100644 --- a/python/freetoken/models/minimax_m3/config.py +++ b/python/freetoken/models/minimax_m3/config.py @@ -34,9 +34,12 @@ ModelConfig, RotaryConfig, ) +from freetoken.utils import init_logger from .args import load_args +logger = init_logger(__name__) + # the config names only the image token; the reference processor reads these two from the tokenizer's fixed added tokens IMAGE_START_ID = 200029 # ]<]start of image[>[ IMAGE_END_ID = 200030 # ]<]end of image[>[ @@ -94,33 +97,18 @@ def parse_vision_config(hf_config: Any) -> VisionConfig | None: ) -# House rule: an env switch that changes WHAT IS SERVED must leave a server-log -# trace. parse_config runs a few times per process (engine + weight loader), so -# each resolved-mode line is logged once per distinct value. -_LOGGED_MODES: set = set() - - -def _log_mode_once(key: str, message: str) -> None: - if key in _LOGGED_MODES: - return - _LOGGED_MODES.add(key) - from freetoken.utils import init_logger - - init_logger(__name__).info(message) - - def parse_config(hf_config: Any) -> ModelConfig: text = _text_config(hf_config) num_layers = text.num_hidden_layers + # House rule: an env switch that changes WHAT IS SERVED must leave a server-log trace; parse_config runs a few times per process. # Dev/testing only: cap the layer count (e.g. FREETOKEN_M3_MAX_LAYERS=5 -> 3 dense # + 2 MoE layers) so the forward path / KV / offload cache can be exercised without # pinning the full ~230 GB of experts. Unset in normal use (GLM-5.2 precedent). _cap = os.environ.get("FREETOKEN_M3_MAX_LAYERS") if _cap and int(_cap) < num_layers: # a cap >= the real count changes nothing num_layers = int(_cap) - _log_mode_once( - f"cap={num_layers}", + logger.info_once( f"FREETOKEN_M3_MAX_LAYERS: serving a TRUNCATED model " f"({num_layers}/{text.num_hidden_layers} layers) -- dev/testing only, " "outputs are garbage; unset the env for real serving.", @@ -128,8 +116,7 @@ def parse_config(hf_config: Any) -> ModelConfig: sparse_enabled = os.getenv("FREETOKEN_M3_SPARSE", "1") != "0" if not sparse_enabled: - _log_mode_once( - "sparse=0", + logger.info_once( "FREETOKEN_M3_SPARSE=0: serving the DENSE-attention ablation (plain " "MHA pool, no indexer; different math and outputs from the shipped " "block-sparse model).", diff --git a/python/freetoken/models/qwen3_vl/vision.py b/python/freetoken/models/qwen3_vl/vision.py index 7c4e0df291..1b8e4bb88f 100644 --- a/python/freetoken/models/qwen3_vl/vision.py +++ b/python/freetoken/models/qwen3_vl/vision.py @@ -16,7 +16,7 @@ LinearRowParallel, OPList, ) -from freetoken.utils import div_even +from freetoken.utils import div_even, init_logger from freetoken.models.weight_stream import BlockWeightStreamer @@ -25,6 +25,8 @@ from freetoken.message import MMItem from .config import VisionConfig +logger = init_logger(__name__) + class VisionLayerNorm(BaseOP): def __init__(self, size: int, eps: float = 1e-6): @@ -42,14 +44,10 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x[..., half:], x[..., :half]), dim=-1) -_warned_rope_fallback = False - - def _apply_vision_rope( q: torch.Tensor, k: torch.Tensor, cache: torch.Tensor, positions: torch.Tensor, head_dim: int ) -> None: """In-place 2D rope through the NeoX text kernel: the [c,c]/[s,s] layout is the half-rotation layout and positions index per-token cache rows.""" - global _warned_rope_fallback S = q.shape[0] try: from freetoken.kernel.triton.rope import apply_rope_with_cos_sin_cache_inplace @@ -60,11 +58,7 @@ def _apply_vision_rope( return except ImportError: # only a missing triton falls back; a CUDA error must not re-run eager on a half-rotated q/k - if not _warned_rope_fallback: - _warned_rope_fallback = True - from freetoken.utils import init_logger - - init_logger(__name__).warning("vision rope triton kernel unavailable; using eager") + logger.warning_once("vision rope triton kernel unavailable; using eager") half = cache.shape[1] // 2 cos = torch.cat((cache[:, :half], cache[:, :half]), dim=-1).unsqueeze(-2) sin = torch.cat((cache[:, half:], cache[:, half:]), dim=-1).unsqueeze(-2) diff --git a/python/freetoken/utils/logger.py b/python/freetoken/utils/logger.py index e5918642be..c2cd5ebb69 100644 --- a/python/freetoken/utils/logger.py +++ b/python/freetoken/utils/logger.py @@ -4,6 +4,8 @@ from typing import TYPE_CHECKING _LOG_LEVEL = None +# module-level: init_logger runs again for the same name and rebuilds the logger's methods +_LOGGED_ONCE: set[tuple] = set() def init_logger( @@ -110,6 +112,13 @@ def _call_rank0(msg, *args, _which, **kwargs): if tp_info is None or tp_info.is_primary(): getattr(logger, _which)(msg, *args, **kwargs) + def _call_once(msg, *args, _which, **kwargs): + key = (name, _which, msg, args) + if key in _LOGGED_ONCE: + return + _LOGGED_ONCE.add(key) + _call_rank0(msg, *args, _which=_which, **kwargs) + if TYPE_CHECKING: class WrapperLogger(logging.Logger): @@ -119,6 +128,10 @@ def info_rank0(self, msg, *args, **kwargs): ... def warning_rank0(self, msg, *args, **kwargs): ... def debug_rank0(self, msg, *args, **kwargs): ... def critical_rank0(self, msg, *args, **kwargs): ... + # rank 0 only, once per distinct (msg, args) in this process; args must hash by value (not exceptions or tensors) + def debug_once(self, msg, *args, **kwargs): ... + def info_once(self, msg, *args, **kwargs): ... + def warning_once(self, msg, *args, **kwargs): ... return WrapperLogger(name) else: @@ -126,4 +139,7 @@ def critical_rank0(self, msg, *args, **kwargs): ... logger.debug_rank0 = partial(_call_rank0, _which="debug") logger.critical_rank0 = partial(_call_rank0, _which="critical") logger.warning_rank0 = partial(_call_rank0, _which="warning") + logger.debug_once = partial(_call_once, _which="debug") + logger.info_once = partial(_call_once, _which="info") + logger.warning_once = partial(_call_once, _which="warning") return logger