Add RL sequence-packing and vLLM block-size knobs

src/maxtext/configs/post_train/rl.yml

@@ -115,6 +115,7 @@ num_batches: 4
 # and/or async sampling and training.
 train_micro_batch_size: -1
 rollout_micro_batch_size: -1
+max_seq_token_per_tpu: 0  # > 0 packs sequences into rows of this many tokens for training; 0 = one padded row per sequence
 # Keep `num_test_batches` low so that evaluation runs quickly. It can be
 # increased to a max. of 330 (if batch size is 4).
 num_test_batches: 5  # 200
@@ -176,6 +177,7 @@ enable_dp_attention: false
 # Performance tuning for samplers
 max_num_batched_tokens: null
 max_num_seqs: null
+vllm_block_size: null  # KV-cache page size; null lets the backend choose
 # If true, enables asynchronous scheduling in vLLM for faster generation
 async_scheduling: true
 # stop generation when any of these strings is generated

src/maxtext/configs/types.py

@@ -2659,6 +2659,11 @@ class VLLM(BaseModel):
   async_scheduling: bool = Field(False, description="Enable asynchronous scheduling in vLLM.")
   max_num_batched_tokens: Optional[int] = Field(None, description="Max number of batched tokens in vLLM.")
   max_num_seqs: Optional[int] = Field(None, description="Max number of sequences in vLLM.")
+  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.",
+  )
   stop_strings: Optional[list[str]] = Field(None, description="List of stop strings for vLLM decoding.")
   vllm_additional_config: dict[str, Any] = Field(default_factory=dict, description="Additional vLLM config options.")
   vllm_hf_overrides: dict[str, Any] = Field(
@@ -2767,6 +2772,15 @@ class RLDataset(BaseModel):
   train_fraction: float = Field(1.0, gt=0.0, le=1.0, description="Fraction of the dataset to be used for training.")
   train_micro_batch_size: int = Field(-1, description="Micro batch size for training.")
   rollout_micro_batch_size: int = Field(-1, description="Micro batch size for rollout.")
+  max_seq_token_per_tpu: int = Field(
+      0,
+      ge=0,
+      description=(
+          "Token budget per packed training row (Tunix `max_seq_token_per_tpu`). When > 0, rollout sequences are packed "
+          "into rows of this many tokens for the actor/reference passes instead of one padded row per sequence; a maximal "
+          "sequence (max_prefill_predict_length + generation length) must fit in one row. 0 disables packing."
+      ),
+  )
   dataset_processor_path: str = Field(
       "",
       description=(

src/maxtext/integration/tunix/tunix_adapter.py

@@ -109,10 +109,18 @@ def __call__(
       decoder_segment_ids: Optional[Array] = None,
       output_hidden_states: bool = False,  # ignored
       forced_routed_experts: Optional[Array] = None,
+      segment_ids: Optional[Array] = None,
   ) -> Tuple[Array, None]:
     """Forward compatible with Tunix Trainers default loss.
     Returns logits, None.
+
+    `segment_ids` is the name Tunix uses for packed-sequence segment ids: it
+    forwards them only to models whose call signature has a parameter of that
+    exact name. They are MaxText's `decoder_segment_ids`; when both are given,
+    `segment_ids` wins so packed rows keep per-sequence attention isolation.
     """
+    if segment_ids is not None:
+      decoder_segment_ids = segment_ids
     if decoder_segment_ids is None and self._pad_id is not None:
       decoder_segment_ids = (input_tokens != self._pad_id).astype(jnp.int32)
     logits = self.base(

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -436,6 +436,18 @@ def create_rl_components(  # pylint: disable=too-many-positional-arguments
 
   rl_rollout_engine = functools.partial(MaxTextVllmRollout, maxtext_config=trainer_config)
 
+  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
+
   cluster_config = rl_cluster_lib.ClusterConfig(
       role_to_mesh={
           rl_cluster_lib.Role.ACTOR: actor_mesh,
@@ -456,6 +468,7 @@ def create_rl_components(  # pylint: disable=too-many-positional-arguments
           mini_batch_size=trainer_config.batch_size,
           train_micro_batch_size=train_micro_batch_size,
           rollout_micro_batch_size=rollout_micro_batch_size,
+          max_seq_token_per_tpu=trainer_config.max_seq_token_per_tpu or None,
           metrics_logging_options=metrics_logging_options,
           profiler_options=profiler_options,
           checkpoint_root_directory=checkpoint_dir,
@@ -484,13 +497,7 @@ def create_rl_components(  # pylint: disable=too-many-positional-arguments
           rollout_vllm_async_scheduling=trainer_config.async_scheduling,
           rollout_vllm_server_mode=trainer_config.rl.use_agentic_rollout,
           rollout_vllm_reshard_chunk_size=trainer_config.rl.reshard_chunk_size,
-          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,
-          },
+          rollout_vllm_kwargs=rollout_vllm_kwargs,
           rollout_vllm_sampling_kwargs={
               "stop": trainer_config.stop_strings,
               "detokenize": trainer_config.stop_strings is not None,

tests/post_training/unit/train_rl_test.py

@@ -53,6 +53,17 @@ def test_rl_config_includes_shared_decoder_defaults(self):
     self.assertFalse(config.enable_mhc_lite)
     self.assertFalse(config.enable_prefix_caching)
 
+  def test_rl_config_packing_and_block_size_defaults(self):
+    """Packing is off and the vLLM page size is backend-chosen unless set."""
+    config = types.RLConfig(model_name="gemma4-26b")
+
+    self.assertEqual(config.max_seq_token_per_tpu, 0)
+    self.assertIsNone(config.vllm_block_size)
+
+    config = types.RLConfig(model_name="gemma4-26b", max_seq_token_per_tpu=12288, vllm_block_size=128)
+    self.assertEqual(config.max_seq_token_per_tpu, 12288)
+    self.assertEqual(config.vllm_block_size, 128)
+
   @pytest.mark.cpu_only
   def test_rollout_prefix_caching_respects_config_for_attention_model(self):
     config = SimpleNamespace(

tests/post_training/unit/tunix_adapter_test.py

@@ -116,6 +116,32 @@ def test_passes_through_explicit_segment_ids_unchanged(self):
 
     np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(explicit_seg))
 
+  def test_segment_ids_alias_maps_to_decoder_segment_ids(self):
+    """Tunix passes packed segment ids under the name `segment_ids`; the
+    adapter must forward them as MaxText's `decoder_segment_ids` instead of
+    synthesizing a pad mask."""
+    adapter = TunixMaxTextAdapter(base_model=self.base, pad_id=99)
+
+    input_tokens = jnp.array([[10, 11, 12, 20, 21]], dtype=jnp.int32)
+    positions = jnp.array([[0, 1, 2, 0, 1]], dtype=jnp.int32)
+    packed_seg = jnp.array([[1, 1, 1, 2, 2]], dtype=jnp.int32)
+
+    adapter(input_tokens, positions, None, None, segment_ids=packed_seg)
+
+    np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(packed_seg))
+
+  def test_segment_ids_alias_takes_precedence_over_decoder_segment_ids(self):
+    adapter = TunixMaxTextAdapter(base_model=self.base, pad_id=99)
+
+    input_tokens = jnp.array([[10, 11, 12, 20, 21]], dtype=jnp.int32)
+    positions = jnp.array([[0, 1, 2, 0, 1]], dtype=jnp.int32)
+    packed_seg = jnp.array([[1, 1, 1, 2, 2]], dtype=jnp.int32)
+    other_seg = jnp.array([[7, 7, 7, 7, 7]], dtype=jnp.int32)
+
+    adapter(input_tokens, positions, None, None, decoder_segment_ids=other_seg, segment_ids=packed_seg)
+
+    np.testing.assert_array_equal(np.asarray(self.base.captured["decoder_segment_ids"]), np.asarray(packed_seg))
+
   def test_forwards_forced_routed_experts(self):
     """The stub captures the kwarg but nothing asserted on it, so deleting the