feat: support RAO grouped agent rollouts

areal/engine/fsdp_engine.py

@@ -1033,6 +1033,12 @@ def _create_device_model(self):
 
         self.get_device_stats().log("before model creation/loading")
 
+        # FA4 不是 transformers 自带的 backend，用到时先注册进 ALL_ATTENTION_FUNCTIONS。
+        # Blackwell(sm100) 上 FA2/FA3 都不可用，只有 FA4 能跑。
+        from areal.engine.fsdp_utils.fa4_attn import maybe_register_fa4
+
+        maybe_register_fa4(self.config.attn_impl)
+
         # Note: VLMs often have vision_tower in fp32 already; loading whole
         # model in optimizer_dtype (fp32 default) is consistent.
         if self.is_vision_model:

areal/engine/fsdp_utils/attn_impl.py

@@ -7,6 +7,9 @@
     "sdpa",
     "flash_attention_2",
     "flash_attention_3",
+    # FA4（flash_attn.cute）。transformers 没有这个 backend，由 fa4_attn.py 注册；
+    # 名字不能含 "flash"（否则被 transformers 抢走）。Blackwell(sm100) 上唯一可用的 FA。
+    "fa4",
     "flex_attention",
 )
 

areal/engine/fsdp_utils/fa4_attn.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+"""把 FlashAttention-4（`flash_attn.cute`）注册成 transformers 的 attention backend。
+
+**为什么需要这个文件**
+
+Blackwell（B200 / sm_100）上：
+  - FA2 没有 sm100 kernel，编不出来；
+  - FA3 是 Hopper-only；
+  - 只有 FA4（`flash-attn-4` 包，import 路径 `flash_attn.cute`）支持 sm100，
+    且前向/反向都有。
+
+但 transformers（截至 5.3.0）没有 `flash_attention_4` 这个 backend——
+它的 `flash_attention_2` 路径走的是 FA2 的老 API `from flash_attn import flash_attn_func`，
+而 `flash-attn-4` 装出来的 `flash_attn` 是个只含 `cute` 子模块的 namespace package，
+于是报 `cannot import name 'flash_attn_func' from 'flash_attn' (unknown location)`。
+参见 huggingface/transformers#44559。
+
+本模块用 transformers 的 `AttentionInterface.register()` 补上这个 backend，
+之后 `attn_implementation="flash_attention_4"` 即可正常使用。
+
+AReaL 侧的输入是**打包过的变长序列**（batch=1，附带 `cu_seq_lens_q/k`、
+`max_length_q/k`），正好对应 FA4 的 `flash_attn_varlen_func`。
+"""
+
+from typing import Any
+
+import torch
+
+from areal.utils import logging
+
+logger = logging.getLogger("FA4")
+
+ATTN_IMPL_NAME = "fa4"  # 名字不能含 "flash"：transformers 的 is_flash_attention_requested 是 `"flash" in name`，
+# 一旦命中就会去 lazy-import FA2/FA3/hub-kernels，绕不过我们注册的实现。
+
+_registered = False
+
+
+def fa4_available() -> bool:
+    """FA4 是否可用（装了 flash-attn-4 且能 import 到 cute 后端）。"""
+    try:
+        from flash_attn.cute import flash_attn_func  # noqa: F401
+        from flash_attn.cute import flash_attn_varlen_func  # noqa: F401
+
+        return True
+    except Exception:
+        return False
+
+
+def _unwrap(out: Any) -> torch.Tensor:
+    """FA4 在 return_lse=True 时返回 (out, lse)，统一取出 out。"""
+    return out[0] if isinstance(out, tuple) else out
+
+
+def fa4_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    sliding_window: int | None = None,
+    softcap: float | None = None,
+    is_causal: bool | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """transformers attention-interface 约定的签名。
+
+    入参 q/k/v 形状是 (batch, num_heads, seq, head_dim)；
+    FA4 和 FA2 一样吃 (…, seq, num_heads, head_dim)，所以先 transpose。
+    """
+    # ── 树训练分支 ─────────────────────────────────────────────────────
+    # enable_tree_training 打开时，AReaL 会把 trie 打包后的注意力掩码作为
+    # tree_block_mask / tree_triton_data 透传进来。它官方的做法是 monkey-patch
+    # transformers.integrations.flash_attention._flash_attention_forward
+    # （module_fsdp.patch_fsdp_for_tree_training），但那条路径只在 attn_impl 是
+    # flash_attention_2 时才被调用 —— 而 B200 上 FA2 根本导入不了
+    # （`cannot import name 'flash_attn_func'`），我们走的是自注册的 fa4。
+    # 不在这里接管的话，树打包后的序列会被当成普通因果序列算，**分支之间会互相
+    # 看见**，是静默的错误结果。
+    if kwargs.get("tree_block_mask") is not None or kwargs.get("tree_triton_data") is not None:
+        from areal.models.tree_attn.module_fsdp import _tree_attn_fwd_func
+
+        out = _tree_attn_fwd_func(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attention_mask,
+            scaling,
+            **kwargs,
+        )
+        return out, None
+
+    from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
+
+    if dropout not in (0.0, None):
+        raise ValueError("FlashAttention-4 不支持 attention dropout（dropout != 0）。")
+
+    # (b, h, s, d) -> (b, s, h, d)
+    q = query.transpose(1, 2)
+    k = key.transpose(1, 2)
+    v = value.transpose(1, 2)
+
+    # 精度：FA4 只接受 bf16/fp16。混精下 query 可能是 fp32，跟随 value 的 dtype。
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        target = v.dtype if v.dtype in (torch.bfloat16, torch.float16) else torch.bfloat16
+        q, k, v = q.to(target), k.to(target), v.to(target)
+
+    causal = module.is_causal if is_causal is None else is_causal
+    window = (sliding_window - 1, 0) if sliding_window else (None, None)
+
+    common = dict(
+        softmax_scale=scaling,
+        causal=causal,
+        window_size=window,
+        softcap=0.0 if softcap is None else float(softcap),
+    )
+
+    cu_q = kwargs.get("cu_seq_lens_q")
+    cu_k = kwargs.get("cu_seq_lens_k")
+
+    if cu_q is not None and cu_k is not None:
+        # 打包变长路径：AReaL 的常规路径（batch 恒为 1）
+        b, s = q.shape[0], q.shape[1]
+        out = flash_attn_varlen_func(
+            q.reshape(b * s, *q.shape[2:]),
+            k.reshape(b * s, *k.shape[2:]),
+            v.reshape(b * s, *v.shape[2:]),
+            cu_seqlens_q=cu_q.to(torch.int32),
+            cu_seqlens_k=cu_k.to(torch.int32),
+            max_seqlen_q=int(kwargs.get("max_length_q", s)),
+            max_seqlen_k=int(kwargs.get("max_length_k", s)),
+            **common,
+        )
+        attn_output = _unwrap(out).reshape(b, s, *q.shape[2:])
+    else:
+        # 稠密路径（没打包时，比如某些 eval / VLM 分支）
+        if attention_mask is not None:
+            raise ValueError(
+                "fa4 目前只支持打包变长输入或无 padding 的稠密输入，"
+                "收到了非空 attention_mask。请改用 sdpa。"
+            )
+        attn_output = _unwrap(flash_attn_func(q, k, v, **common))
+
+    return attn_output.to(value.dtype), None
+
+
+def register_fa4() -> bool:
+    """把 FA4 注册进 transformers 的 ALL_ATTENTION_FUNCTIONS。幂等。
+
+    Returns
+    -------
+    bool
+        注册成功（或此前已注册）返回 True；FA4 不可用返回 False。
+    """
+    global _registered
+    if _registered:
+        return True
+    if not fa4_available():
+        return False
+
+    from transformers import AttentionInterface
+
+    AttentionInterface.register(ATTN_IMPL_NAME, fa4_attention_forward)
+    _registered = True
+    logger.info(f"Registered transformers attention backend '{ATTN_IMPL_NAME}'.")
+    return True
+
+
+def maybe_register_fa4(attn_impl: str) -> None:
+    """建模型前调用：若用户选了 flash_attention_4，就确保它已注册。"""
+    if attn_impl != ATTN_IMPL_NAME:
+        return
+    if not register_fa4():
+        raise RuntimeError(
+            "attn_impl='fa4' 但 FA4 不可用。"
+            "请安装 flash-attn-4（`from flash_attn.cute import flash_attn_func` 需能 import）。"
+        )

areal/experimental/openai/cache.py

@@ -263,6 +263,15 @@ def __setitem__(
         value: InteractionWithTokenLogpReward,
     ) -> None:
         """Add a new interaction to the cache, automatically building parent-child relationships."""
+        # ⭐ RAO：把本次请求原始的 `model`（含 `__dN` 深度后缀）钉到 interaction 上。
+        #    proxy 在转发前会丢掉 model 字段，这里是唯一还拿得到它的地方。
+        if getattr(value, "rao_model", None) is None:
+            try:
+                from areal.experimental.openai.types import RAO_REQ_MODEL
+
+                value.rao_model = RAO_REQ_MODEL.get()
+            except Exception:
+                pass
         if value.messages is None:
             raise ValueError(
                 "Interaction messages must be set to find parent relationship."
@@ -468,6 +477,13 @@ def export_interactions(
         if len(complete_cache) == 0:
             return {}
 
+        # 盖上 session 章：一个 InteractionCache 就是一次 agent 运行（一条 rollout
+        # session）。下游 dump 靠它还原"哪几轮属于同一条会话"——不能靠 parent 链，
+        # 因为 harness 丢弃被截断的回复时，parent_data 不再是 child 的前缀，
+        # 父子关系会断（见本文件 _is_similar_on_last_message 的 Prefix mismatch 告警）。
+        for interaction in complete_cache.values():
+            interaction.session_id = self._session_id
+
         if style == "concat":
             for interaction in complete_cache.values():
                 if interaction.chat_template_type != "concat":

areal/experimental/openai/client.py

@@ -357,6 +357,33 @@ def _resolve_max_total_tokens(
     return min(prompt_len + max_new_tokens, cap)
 
 
+def _flatten_text_content_parts(messages: list[dict]) -> list[dict]:
+    """把纯文本的 OpenAI content-parts 列表压回普通字符串。
+
+    有些 harness（如 pi）按 OpenAI 多模态格式发消息，``content`` 是
+    ``[{"type": "text", "text": "..."}]`` 而不是字符串。而不少 chat template
+    对非字符串 content 直接置空 —— Qwen3 的模板就是：
+
+        {%- if message.content is string %}{%- set content = message.content %}
+        {%- else %}{%- set content = '' %}{%- endif %}
+
+    结果是**用户消息被静默渲染成空串**，模型收到一个没有内容的请求，
+    还会一本正经地回"你还没给我任务"。这里在套模板之前压平，只动纯文本的情况；
+    含图片等其它 part 的消息原样保留，交给 _extract_images_from_messages。
+    """
+    out = []
+    for m in messages:
+        c = m.get("content")
+        if (
+            isinstance(c, list)
+            and c
+            and all(isinstance(x, dict) and x.get("type") == "text" for x in c)
+        ):
+            m = {**m, "content": "".join(x.get("text") or "" for x in c)}
+        out.append(m)
+    return out
+
+
 def _parse_tool_call_arguments(messages: list[dict]) -> list[dict]:
     """Return a new message list with tool_call arguments parsed from JSON strings
     to dicts. Some chat templates (e.g. GLM-5.1) iterate over arguments with
@@ -662,6 +689,7 @@ async def create(
         has_images = len(image_data) > 0
 
         tokenizer_messages = messages_for_tokenizer if has_images else messages_list
+        tokenizer_messages = _flatten_text_content_parts(tokenizer_messages)
         tokenizer_messages = _parse_tool_call_arguments(tokenizer_messages)
         if self.chat_template_type == "hf":
             prompt_token_ids = apply_chat_template(
@@ -684,6 +712,10 @@ async def create(
                 )
             else:
                 concat_tok_messages = concat_messages
+            # ★ concat 分支用的是 remaining_messages，绕开了上面那次扁平化，
+            #   这里必须再压一次，否则 pi 这类发 content-parts 的 harness
+            #   会被模板渲染成空消息。
+            concat_tok_messages = _flatten_text_content_parts(concat_tok_messages)
             prompt_token_ids = concat_prompt_token_ids_with_parent(
                 concat_tok_messages,
                 interaction.parent if interaction is not None else None,

areal/experimental/openai/proxy/proxy_rollout_server.py

@@ -30,6 +30,7 @@
 
 from areal.api.cli_args import NameResolveConfig
 from areal.experimental.openai.client import ArealOpenAI
+from areal.experimental.openai.types import RAO_REQ_MODEL
 from areal.infra.rpc.serialization import deserialize_value, serialize_value
 from areal.infra.utils.http import validate_admin_api_key
 from areal.utils import name_resolve, names, seeding
@@ -578,6 +579,21 @@ async def _call_client_create(
     session_data.update_last_access()
 
     sig = inspect.signature(create_fn)
+    # ⭐ RAO：`model` 马上就要被丢掉（AReaL 的 client 只服务单一模型、不接受该参数），
+    #    但 harness 把**递归深度**编码在 model id 尾部（`<id>__dN`）。丢掉之后
+    #    深度信息就永久消失了 —— 实测表现是训练数据里每条序列的 rao_depth 都是 0。
+    #    所以在丢弃之前先存进 ContextVar，interaction 落库时抄到 `rao_model` 上。
+    #    用 ContextVar 而非实例属性：同一 session 的多个 sub-agent 是并行请求的。
+    try:
+        _req_model = (
+            request.get("model")
+            if isinstance(request, dict)
+            else getattr(request, "model", None)
+        )
+        if _req_model:
+            RAO_REQ_MODEL.set(str(_req_model))
+    except Exception:
+        pass  # 观测失败绝不能影响请求
     areal_client_ignored_args = ["model"] + (extra_ignored_args or [])
     areal_client_disallowed_args = ["areal_cache"]
     areal_client_allowed_args = list(

areal/experimental/openai/proxy/server.py

@@ -140,18 +140,30 @@ def serialize_interactions(
 
     result = {}
     for key, interaction in interactions.items():
+        # 会话身份必须一起过河：subproc/online 模式下 interaction 要跨进程传回，
+        # 不带上的话下游只能看到一堆扁平的 turn，无法还原"哪几轮属于同一条会话"。
+        identity = {
+            "session_id": interaction.session_id,
+            "parent_interaction_id": (
+                interaction.parent.interaction_id
+                if interaction.parent is not None
+                else interaction.parent_interaction_id
+            ),
+        }
         if interaction.has_tensor_data:
             result[key] = {
                 "tensor_dict": interaction.to_tensor_dict(),
                 "reward": interaction.reward,
                 "interaction_id": interaction.interaction_id,
+                **identity,
             }
         else:
             result[key] = {
                 "messages": interaction.messages,
                 "output_message_list": interaction.output_message_list,
                 "reward": interaction.reward,
                 "interaction_id": interaction.interaction_id,
+                **identity,
             }
     return serialize_value(result)
 
@@ -174,6 +186,8 @@ def deserialize_interactions(
             interaction.output_message_list = item["output_message_list"]
         interaction.reward = item["reward"]
         interaction.interaction_id = item["interaction_id"]
+        interaction.session_id = item.get("session_id")
+        interaction.parent_interaction_id = item.get("parent_interaction_id")
         result[key] = interaction
     return result