Add RL sequence-packing and vLLM block-size knobs

[gutianyu-google]

Description

Tunix packs RL training sequences into fixed-token rows when its training config carries max_seq_token_per_tpu, and it forwards the packed segment ids to the model only when the model's call signature has a parameter literally named segment_ids. MaxText exposed neither, so packing could not be turned on from an RL config: the CLI rejected the key (not in RLConfig), and the adapter would never have received the packed ids.

This PR adds:

• max_seq_token_per_tpu to RLConfig (default 0 = unpacked, one padded row per sequence, i.e. today's behavior), plumbed into RLTrainingConfig. When set, a maximal sequence (max_prefill_predict_length + generation length) must fit in one row; Tunix validates this at startup.
• A segment_ids parameter on TunixMaxTextAdapter.__call__, mapped to MaxText's decoder_segment_ids, so packed rows keep per-sequence attention isolation. Without it the adapter falls back to synthesizing a pad mask and packed sequences would attend to each other.
• vllm_block_size to RLConfig (default None, backend-chosen as today). tpu_inference derives the KV-cache page size from the engine shape, so unrelated engine changes (max_model_len, TP/DP layout) silently move it; pinning it keeps generation comparable across such experiments. It is forwarded through rollout_vllm_kwargs only when set.

Measured on a Qwen3-0.6B GRPO run on TPU v7x (2048 sequences/step, 12288-token rows, prompts ~150 tokens, completions ~4400 tokens): packing cut actor_train_time from 15.2 s to 6.7 s per step, with rewards and completion lengths unchanged. Both defaults leave existing configs untouched.

Tests

• tests/post_training/unit/tunix_adapter_test.py: segment_ids is forwarded as decoder_segment_ids, and takes precedence when both are given.
• tests/post_training/unit/train_rl_test.py: RLConfig defaults and explicit values for both knobs.

python -m pytest tests/post_training/unit/tunix_adapter_test.py tests/post_training/unit/train_rl_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces support for sequence packing and custom KV-cache block sizes in RL post-training by adding max_seq_token_per_tpu and vllm_block_size configurations, updating the Tunix adapter to handle packed-sequence segment IDs, and adding corresponding unit tests. The reviewer feedback suggests adding validation to ensure vllm_block_size is a power of 2 and verifying that max_seq_token_per_tpu is at least max_target_length when sequence packing is enabled to prevent runtime failures.

[gemini-code-assist]

vLLM block sizes must be powers of 2 (typically 8, 16, 32, 64, 128, 256). Adding a field validator ensures that any invalid non-power-of-2 block size is caught early during configuration validation rather than causing a startup failure in vLLM.

Suggested change

	vllm_block_size: Optional[int] = Field(
	None,
	gt=0,
	description="KV-cache block (page) size for vLLM. None lets the backend pick it from the engine shape.",
	)
	vllm_block_size: Optional[int] = Field(
	None,
	gt=0,
	description="KV-cache block (page) size for vLLM. None lets the backend pick it from the engine shape.",
	)
	@field_validator("vllm_block_size")
	@classmethod
	def validate_vllm_block_size(cls, v: Optional[int]) -> Optional[int]:
	if v is not None and (v & (v - 1)) != 0:
	raise ValueError("vllm_block_size must be a power of 2.")
	return v

[gemini-code-assist]

When sequence packing is enabled (max_seq_token_per_tpu > 0), a maximal sequence of length max_target_length must fit in a single row. If max_seq_token_per_tpu is configured to be less than max_target_length, Tunix will fail at startup. Adding an early validation check prevents this runtime failure.

  if trainer_config.max_seq_token_per_tpu > 0 and trainer_config.max_seq_token_per_tpu < trainer_config.max_target_length:
    raise ValueError(
        f"max_seq_token_per_tpu ({trainer_config.max_seq_token_per_tpu}) must be greater than or equal to "
        f"max_target_length ({trainer_config.max_target_length}) when sequence packing is enabled."
    )

  rollout_vllm_kwargs = {
      "hf_overrides": trainer_config.vllm_hf_overrides,
      "enable_expert_parallel": sampler_config.enable_expert_parallel,
      "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config),
      # Ensures vLLM model initializes with correct dtype (not float32 default)
      "dtype": trainer_config.weight_dtype.value,
  }
  if trainer_config.vllm_block_size is not None:
    # Pin the KV-cache page size; left unset, the backend derives it from the
    # engine shape, so unrelated engine changes can move it.
    rollout_vllm_kwargs["block_size"] = trainer_config.vllm_block_size