Add exact token content attribution

docs/docs.json

@@ -67,6 +67,7 @@
           "features/checkpoint-forking",
           "features/checkpoint-deletion",
           "features/additional-histories",
+          "features/token-content-attribution",
           "features/tracking-metrics",
           "features/mcp-rl"
         ]

docs/features/token-content-attribution.mdx

@@ -0,0 +1,65 @@
+---
+title: "Token content attribution"
+description: "Distinguish message-body tokens from chat-template scaffolding."
+---
+
+`art.TokenFlag.CONTENT` identifies tokens that came from message bodies:
+caller-provided text, reasoning, tool-call data, or a model-sampled assistant
+emission. Renderer-added role tags, tool declarations, wrappers, separators,
+and generation prompts are scaffold and do not carry the flag. A sampled
+turn-close token is CONTENT; a separator added later while rendering the next
+prompt is not.
+
+At a text boundary, a token belongs to the segment containing its first source
+character. A character exactly on a boundary belongs to the later segment.
+This is the policy implemented by Prime Intellect `renderers==0.1.8`, whose
+`RenderedTokens.is_content` signal ART reuses.
+
+## Availability and exactness
+
+Inspect CONTENT only when `TokenizedTrajectory.content_attribution_complete`
+is `True`. ART sets it only when every prompt in the sequence has exact token
+IDs from the inference engine and the selected renderer matches each complete
+prompt ID-for-ID. If attribution is unavailable, CONTENT is unset on every
+token; that means unknown, not known scaffold.
+
+The renderer-qualified Megatron matrix is:
+
+| Model family | Qualification |
+| --- | --- |
+| Qwen3 dense and MoE chat models in ART's registry | Text, reasoning, tools, and multi-turn histories; exact parity is checked again per prompt. |
+| Qwen3.5 and Qwen3.6 dense and MoE models in ART's registry | Text, reasoning, tools, and multi-turn histories; both engine tool-argument representations are accepted only by exact match. |
+| `OpenPipe/Qwen3-14B-Instruct` | Text and tool declarations. Multi-turn assistant tool-call histories currently fail closed. |
+| `openai/gpt-oss-20b` and `openai/gpt-oss-120b` | Canonical `openai-harmony` format, including reasoning and function tools. |
+| Gemma 4 IT and DeepSeek V4 Flash/Pro | Unavailable: no exact renderer/parity path has been established. |
+
+Explicit base models, raw Completions, custom chat templates, and models outside
+that matrix also fail closed. The Prime renderer dependency is installed by
+ART's `backend` and `megatron` extras; a core-only installation reports content
+attribution as unavailable.
+
+Chat Completions, Anthropic Messages, and OpenAI Responses use the same rule.
+Their final response (including a reconstructed streamed response) must carry
+the exact `prompt_token_ids` consumed by the engine. Caladan's compatible
+vLLM/SGLang protocol patches supply this metadata; other servers qualify only
+if they provide the same exact carrier. A Responses request that performs more
+than one internal server-side generation does not have one unambiguous prompt
+sequence and therefore omits attribution.
+
+## Selecting losses safely
+
+Require both complete attribution and at least one selected token:
+
+```python
+if not tokenized.content_attribution_complete:
+    raise ValueError("CONTENT attribution is unavailable")
+
+content_mask = [bool(flag & art.TokenFlag.CONTENT) for flag in tokenized.flags]
+if not any(content_mask):
+    raise ValueError("CONTENT selection is empty")
+
+mean_content_loss = losses[content_mask].mean()
+```
+
+The second check is required because NumPy and PyTorch return NaN for
+`losses[empty_mask].mean()`.

pyproject.toml

@@ -45,6 +45,7 @@ backend = [
     "gql>=4.0.0",
     "nvidia-cudnn-frontend<1.21 ; sys_platform == 'linux'",
     "nvidia-resiliency-ext<0.5 ; sys_platform == 'linux'",
+    "renderers==0.1.8",
 ]
 megatron = [
     "numpy<2",
@@ -73,6 +74,7 @@ megatron = [
     "transformers==5.12.1",
     "scipy>=1.17.0",
     "ml-dtypes>=0.5.0 ; python_full_version < '3.13'",
+    "renderers==0.1.8",
 ]
 
 langgraph = [
@@ -332,6 +334,7 @@ dev = [
     "prek>=0.2.29",
     "uv>=0.11.7",
     "skypilot[cudo,do,fluidstack,gcp,kubernetes,lambda,paperspace,runpod]==0.11.1",
+    "renderers==0.1.8",
 ]
 
 [tool.uv.sources]

src/art/local/backend.py

@@ -208,6 +208,7 @@ def _apply_configured_chat_template(
     chat_template = _configured_chat_template_value(internal_config)
     if chat_template is not None:
         tokenizer.chat_template = chat_template
+        setattr(tokenizer, "_art_custom_chat_template", True)
 
 
 def _model_support_handler(

src/art/trajectories/__init__.py

@@ -64,6 +64,14 @@ class TokenFlag(IntFlag):
     EXACT = 1 << 0
     # The token belongs to a model-sampled response rather than its prompt.
     SAMPLED = 1 << 1
+    # The token comes from a message body: caller-provided content, reasoning,
+    # or tool-call data, or a model-sampled assistant emission. Renderer-added
+    # role tags, tools headers, separators, and generation prompts are not
+    # content. For tokens spanning a body/scaffold text boundary, the segment
+    # containing the token's first source character owns the token; a character
+    # exactly on a boundary belongs to the later segment. This bit is meaningful
+    # only when TokenizedTrajectory.content_attribution_complete is true.
+    CONTENT = 1 << 2
 
 
 class ChatCompletionsRequest(TypedDict, total=False, extra_items=Any):
@@ -529,9 +537,21 @@ def __len__(self) -> int:
 
 
 class TokenizedTrajectory(pydantic.BaseModel):
+    """One token sequence with independent provenance flags.
+
+    Consumers selecting CONTENT tokens must first require
+    ``content_attribution_complete``, then require a nonempty selection. In
+    NumPy and PyTorch, ``losses[empty_mask].mean()`` is NaN; silently averaging
+    an unavailable or empty CONTENT mask can therefore poison a training step.
+    """
+
     token_ids: list[int]
     logprobs: list[float]
     flags: list[TokenFlag]
+    # True only when every token has a known CONTENT classification, proven by
+    # matching renderer output to every exact engine-consumed prompt in full.
+    # When false, an unset CONTENT bit is unknown, not known scaffold.
+    content_attribution_complete: bool = False
     underlying: Trajectory
 
 

src/art/trajectories/_content.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+import copy
+import json
+import logging
+from typing import Any, Protocol, cast
+
+logger = logging.getLogger(__name__)
+
+
+class _RenderedTokens(Protocol):
+    token_ids: list[int]
+    is_content: list[bool]
+
+
+class _Renderer(Protocol):
+    def render(
+        self,
+        messages: list[dict[str, Any]],
+        *,
+        tools: object,
+        add_generation_prompt: bool,
+    ) -> _RenderedTokens: ...
+
+
+def _renderer_name(base_model: str) -> str | None:
+    # Importing the registry is safe here: its model sets do not import Megatron
+    # or model handlers. Keeping the allowlist there prevents CONTENT support
+    # from silently expanding to unqualified fine-tunes with different templates.
+    from art.megatron.model_support.registry import (
+        GPT_OSS_MOE_MODELS,
+        QWEN3_5_MODELS,
+        QWEN3_DENSE_MODELS,
+        QWEN3_MOE_MODELS,
+    )
+
+    if base_model.endswith("-Base"):
+        return None
+    if base_model in QWEN3_DENSE_MODELS | QWEN3_MOE_MODELS:
+        return "qwen3"
+    if base_model in QWEN3_5_MODELS:
+        return "qwen3.6" if "/Qwen3.6-" in base_model else "qwen3.5"
+    if base_model in GPT_OSS_MOE_MODELS:
+        return "gpt-oss"
+    return None
+
+
+def _renderer_messages(
+    messages: list[dict[str, Any]], *, parse_tool_arguments: bool
+) -> list[dict[str, Any]]:
+    normalized = copy.deepcopy(messages)
+    for message in normalized:
+        reasoning = message.pop("reasoning", None)
+        if reasoning is not None and "reasoning_content" not in message:
+            message["reasoning_content"] = reasoning
+        refusal = message.get("refusal")
+        if isinstance(refusal, str) and refusal and not message.get("content"):
+            message["content"] = refusal
+        for tool_call in message.get("tool_calls") or []:
+            if not isinstance(tool_call, dict):
+                continue
+            function = tool_call.get("function")
+            if not isinstance(function, dict):
+                continue
+            arguments = function.get("arguments")
+            if not parse_tool_arguments or not isinstance(arguments, str):
+                continue
+            try:
+                parsed = json.loads(arguments)
+            except json.JSONDecodeError:
+                continue
+            if isinstance(parsed, dict):
+                function["arguments"] = parsed
+    return normalized
+
+
+def _renderer_tool_variants(tools: object) -> list[object]:
+    if not isinstance(tools, list):
+        return [tools]
+    sglang = copy.deepcopy(tools)
+    vllm: list[object] = []
+    changed = False
+    for original, strict_tool in zip(tools, sglang, strict=True):
+        data = original if isinstance(original, dict) else None
+        strict_data = strict_tool if isinstance(strict_tool, dict) else None
+        function = data.get("function") if data is not None else None
+        strict_function = (
+            strict_data.get("function") if strict_data is not None else None
+        )
+        if not isinstance(function, dict) or not isinstance(strict_function, dict):
+            vllm.append(copy.deepcopy(original))
+            continue
+        function = cast(dict[str, Any], function)
+        strict_function = cast(dict[str, Any], strict_function)
+        changed = True
+        strict_function.setdefault("strict", False)
+        vllm.append(
+            {
+                "type": "function",
+                "function": {
+                    "name": function.get("name"),
+                    "parameters": function.get("parameters", {}),
+                    "strict": function.get("strict"),
+                    "type": "function",
+                    "allowed_callers": function.get("allowed_callers"),
+                    "defer_loading": function.get("defer_loading"),
+                    "description": function.get("description"),
+                    "output_schema": function.get("output_schema"),
+                },
+            }
+        )
+    return [tools, sglang, vllm] if changed else [tools]
+
+
+def content_mask_for_exact_prompt(
+    *,
+    tokenizer: object,
+    base_model: str,
+    messages: list[dict[str, Any]],
+    tools: object,
+    expected_prompt_ids: list[int],
+    chat_template_kwargs: Mapping[str, object] | None,
+    has_custom_chat_template: bool,
+) -> list[bool] | None:
+    """Return complete body attribution only after exact prompt-ID parity.
+
+    Prime Intellect's hand-coded renderers provide the body/scaffold boundary
+    and its documented first-source-character policy. ART accepts that signal
+    only for explicitly qualified Megatron model names, without a custom chat
+    template, and only when the renderer's whole prompt exactly equals the IDs
+    consumed by the inference engine. Any uncertainty returns ``None``.
+    """
+
+    renderer_name = _renderer_name(base_model)
+    if renderer_name is None or has_custom_chat_template:
+        return None
+    try:
+        from renderers import config_from_name, create_renderer
+        from renderers.configs import Qwen3RendererConfig
+
+        config = config_from_name(renderer_name)
+        if base_model == "OpenPipe/Qwen3-14B-Instruct":
+            config = Qwen3RendererConfig(enable_thinking=False)
+        config_fields = getattr(type(config), "model_fields", {})
+        ignored_standard_kwargs = (
+            {"enable_thinking", "thinking_budget"}
+            if renderer_name == "gpt-oss"
+            else {"thinking_budget"}
+        )
+        unknown_kwargs = set(chat_template_kwargs or {}) - set(config_fields)
+        if not unknown_kwargs <= ignored_standard_kwargs:
+            return None
+        renderer_kwargs = {
+            key: value
+            for key, value in (chat_template_kwargs or {}).items()
+            if key in config_fields
+        }
+        renderer = cast(
+            _Renderer,
+            create_renderer(
+                tokenizer,
+                cast(Any, config),
+                chat_template_kwargs=renderer_kwargs,
+            ),
+        )
+    except (ImportError, ModuleNotFoundError):
+        return None
+    except Exception:
+        logger.debug(
+            "CONTENT renderer rejected prompt for %s",
+            base_model,
+            exc_info=True,
+        )
+        return None
+
+    # vLLM currently preserves JSON-string tool arguments for templates such
+    # as Qwen3, while SGLang normalizes them to objects for templates such as
+    # Qwen3.5. Both are engine-consumed forms. Try the two renderer inputs and
+    # accept only the one whose entire token sequence proves exact parity.
+    for tool_variant in _renderer_tool_variants(tools):
+        for parse_tool_arguments in (False, True):
+            try:
+                rendered = renderer.render(
+                    _renderer_messages(
+                        messages, parse_tool_arguments=parse_tool_arguments
+                    ),
+                    tools=tool_variant,
+                    add_generation_prompt=True,
+                )
+            except Exception:
+                logger.debug(
+                    "CONTENT renderer rejected a prompt variant for %s",
+                    base_model,
+                    exc_info=True,
+                )
+                continue
+            token_ids = list(rendered.token_ids)
+            is_content = list(rendered.is_content)
+            if token_ids != expected_prompt_ids:
+                continue
+            if len(is_content) != len(token_ids) or not all(
+                isinstance(value, bool) for value in is_content
+            ):
+                continue
+            return is_content
+    return None