[RFC]: Rebase Additional Information into Model Intermediate Buffer

Motivation

The current design uses multiple overlapping concepts:

• additional_information (request-level payload and connector payload)
• runtime_additional_information (model runner runtime kwargs)
• prompt-embed side paths that are structurally similar to additional information

This creates ambiguity in ownership, naming, and lifecycle. It also makes transfer behavior implicit in code instead of explicit in model config.

This RFC proposes a clean rebase:

1. Remove additional_information.
2. Rename runtime_additional_information to model_intermediate_buffer.
3. Remove additional_information payload structs.
4. Remove prompt-embed-related mechanisms from the old additional-information design.
5. Transfer hidden states and ids through connector as dict[str, Tensor].
6. Define transfer scope in model config.

Proposed Change.

Goals

• Establish a single intermediate data concept: model_intermediate_buffer.
• Standardize connector transfer payload as dict[str, Tensor].
• Make transfer schema explicit and configurable in model config.
• Align key naming with connector naming convention.
• Define model-runner internal storage as a request-indexed dictionary.

Non-Goals

• No change to connector backend implementation (shared memory/msgpack internals can remain).
• No algorithmic change to model generation logic.
• No long-term dual-write support beyond migration window.

Key Design

1) Remove additional_information

Delete all request/scheduler/model-runner/public API usage of:

• additional_information
• AdditionalInformationPayload
• AdditionalInformationEntry

2) Rename runtime_additional_information to model_intermediate_buffer

All runtime kwargs and storage in model runner should use:

• model_intermediate_buffer

In the model runner, model_intermediate_buffer is a dictionary:

• Type: dict[str, dict[str, Tensor]]
• Structure: {request_id: intermediate_results}
• intermediate_results type: dict[str, Tensor]
• Example:
  • model_intermediate_buffer["req_abc"] = {"req_abc-24": hs_tensor, "req_abc-thinker_input_ids": ids_tensor}

Examples:

• remove per-request req_state.additional_information_cpu
• introduce centralized runner-level self.model_intermediate_buffer
• _gather_runtime_additional_information() -> _gather_model_intermediate_buffer_map()
• model forward kwarg:
  • before: runtime_additional_information
  • after: model_intermediate_buffer

3) Connector transfer format

Connector payload is strictly:

• dict[str, Tensor]

No custom payload wrapper for additional information is kept.

4) Connector-compatible key naming

The key name must match connector naming style and be deterministic.

Hidden states key

• Format: rid-layer_id
• Example: req_abc-24

IDs key

• Format: rid-id_type
• Example: req_abc-thinker_input_ids
• Example: req_abc-code_predictor_codes

rid is the request id used in connector transfer context.

5) Remove prompt-embed related legacy path

Remove prompt-embed paths that were parallel to old additional-information flow, including:

• prompt-embed payload wrappers in omni request path
• prefill overlay logic that depends on prompt-embed additional path

Config Design (Model Config Driven Transfer)

Transfer scope must be explicitly defined in model config.

engine_args:
  transfer_config:
    enabled: true
    # Which categories are allowed to transfer.
    items:
      - hidden_states
      - ids

    # Key naming templates must match connector-compatible style.
    # hidden_states => rid-layer_id
    # ids => rid-id_type
    key_templates:
      hidden_states: "{rid}-{layer_id}"
      ids: "{rid}-{id_type}"

    # Optional filters to reduce transfer volume.
    hidden_state_layers: [0, 24]
    id_types:
      - thinker_input_ids
      - code_predictor_codes

Config Rules

• items controls which information categories are transferable.
• key_templates is mandatory when enabled=true.
• Hidden states must produce keys as rid-layer_id.
• IDs must produce keys as rid-id_type.
• Any key not matching template should be rejected with clear error logs.

Data Flow (After Rebase)

1. Stage output is postprocessed into connector payload dict[str, Tensor].
2. Connector sends/receives tensors keyed by:
   • hidden states: rid-layer_id
   • ids: rid-id_type
3. Model runner stores data in a centralized map:
   • model_intermediate_buffer: dict[str, dict[str, Tensor]]
   • keyed by request id, value is per-request intermediate tensor map
4. Model forward receives model_intermediate_buffer via kwargs.
5. Model preprocess/postprocess updates merge back into model_intermediate_buffer.

Migration Plan

Phase 1 (Compatibility)

• Accept both old and new runtime kwargs in model forward.
• Internally normalize into model_intermediate_buffer.
• Emit deprecation warning on old names.

Phase 2 (Switch Default)

• Write only new names/paths.
• Keep read compatibility for old keys temporarily.

Phase 3 (Cleanup)

• Remove old symbols and compatibility branches.
• Remove payload structs and prompt-embed legacy path entirely.

Feedback Period.

No response

CC List.

@hsliuustc0106 @Gaohan123 @ywang96 @R2-Y @amy-why-3459

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.