[megatron] fix: support FP8 padding for BSHD actor forward

tests/models/test_mcore_bshd_fp8_padding.py

@@ -0,0 +1,59 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import torch
+
+from verl.models.mcore.util import postprocess_bshd, preprocess_bshd
+
+
+def test_preprocess_bshd_fp8_padding_aligns_local_tokens_and_roundtrips():
+    input_ids = torch.tensor(
+        [
+            [10, 11, 12, 0, 0],
+            [20, 21, 22, 23, 24],
+            [30, 31, 0, 0, 0],
+        ],
+        dtype=torch.long,
+    )
+    attention_mask = torch.tensor(
+        [
+            [True, True, True, False, False],
+            [True, True, True, True, True],
+            [True, True, False, False, False],
+        ]
+    )
+    position_ids = torch.arange(input_ids.shape[1], dtype=torch.long).unsqueeze(0).expand_as(input_ids)
+
+    with (
+        patch("verl.models.mcore.util.mpu.get_context_parallel_world_size", return_value=1),
+        patch("verl.models.mcore.util.mpu.get_tensor_model_parallel_world_size", return_value=2),
+    ):
+        input_ids_bshd, attention_mask_bshd, position_ids_bshd = preprocess_bshd(
+            input_ids,
+            attention_mask,
+            position_ids,
+            sequence_parallel=True,
+            pre_process=True,
+            use_fp8_padding=True,
+        )
+
+    assert input_ids_bshd.shape == (3, 256)
+    assert position_ids_bshd.shape == (3, 256)
+    assert attention_mask_bshd.shape == (3, 256)
+    assert (input_ids_bshd.shape[0] * input_ids_bshd.shape[1] // 2) % 128 == 0
+
+    restored = postprocess_bshd(input_ids_bshd, attention_mask_bshd, attention_mask, input_ids.shape[1])
+    torch.testing.assert_close(restored, input_ids)

verl/models/mcore/model_forward.py

@@ -128,8 +128,6 @@ def model_forward(
             When using the bshd format, we have to add paddings to the input_ids to meet the longest sequence length, 
             so it is recommended to disable dynamic batch size and set batch size to 1
             """
-            assert fp8 is None, "fp8 is not supported for bshd format yet"
-
             batch_size, sequence_length = attention_mask.shape[:2]
             position_ids_for_preprocess = (
                 torch.arange(sequence_length, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
@@ -143,6 +141,7 @@ def model_forward(
                 position_ids_for_preprocess,
                 sequence_parallel=sp,
                 pre_process=pre_process_for_bshd,
+                use_fp8_padding=use_fp8_padding,
             )
             output_orig = model(
                 input_ids=new_input_ids,
@@ -153,7 +152,12 @@ def model_forward(
             if post_process and logits_processor is not None:
                 args = {
                     k: preprocess_bshd(
-                        v, attention_mask, position_ids_for_preprocess, sequence_parallel=sp, pre_process=True
+                        v,
+                        attention_mask,
+                        position_ids_for_preprocess,
+                        sequence_parallel=sp,
+                        pre_process=True,
+                        use_fp8_padding=use_fp8_padding,
                     )[0]
                     for k, v in logits_processor_args.items()
                 }

verl/models/mcore/util.py

@@ -40,6 +40,21 @@ def _compute_fp8_thd_align_size(align_size: int) -> tuple[int, int]:
     return math.lcm(16, align_size), align_size * 128
 
 
+def _align_bshd_max_seqlen_for_fp8(max_seqlen: int, batch_size: int, tp_size: int, cp_size: int = 1) -> int:
+    """Align BSHD sequence length for TE FP8 block quantization.
+
+    In BSHD format, Transformer Engine sees a local activation extent of
+    ``batch_size * max_seqlen / (tp_size * cp_size)``. Block FP8 kernels require
+    that extent to be divisible by 128, while SP/CP still require the sequence
+    length to split cleanly across their ranks.
+    """
+    align_size = tp_size * cp_size * 2 if cp_size > 1 else tp_size
+    fp8_total_align = 128 * tp_size * cp_size
+    fp8_seq_align = fp8_total_align // math.gcd(batch_size, fp8_total_align)
+    fp8_seq_align = math.lcm(fp8_seq_align, align_size)
+    return ((max_seqlen + fp8_seq_align - 1) // fp8_seq_align) * fp8_seq_align
+
+
 def preprocess_packed_seqs(
     input_ids: torch.Tensor, attention_mask: torch.Tensor, pre_process: bool = True, use_fp8_padding: bool = False
 ) -> tuple[torch.Tensor, PackedSeqParams]:
@@ -196,6 +211,7 @@ def preprocess_bshd(
     position_ids: torch.Tensor,
     sequence_parallel: bool = False,
     pre_process: bool = True,
+    use_fp8_padding: bool = False,
 ):
     """
     Remove left padding from input_ids, attention_mask and position_ids
@@ -213,6 +229,14 @@ def preprocess_bshd(
         sp_world_size = mpu.get_tensor_model_parallel_world_size()
         pad_size = (sp_world_size - seq_len % sp_world_size) % sp_world_size
         seq_len = seq_len + pad_size
+    if use_fp8_padding:
+        tp_size = mpu.get_tensor_model_parallel_world_size()
+        seq_len = _align_bshd_max_seqlen_for_fp8(
+            max_seqlen=seq_len,
+            batch_size=batch_size,
+            tp_size=tp_size,
+            cp_size=cp_size,
+        )
     shape[1] = seq_len
     if pre_process:
         new_input_ids = torch.zeros(dtype=input_ids.dtype, device=input_ids.device, size=shape)
@@ -521,12 +545,12 @@ def preprocess_bshd_no_padding(
     if use_fp8_padding:
         # For FP8 block quantization, batch_size * max_seqlen / tp_size must be divisible by 128.
         # We need: max_seqlen % tp_size == 0 (for SP) AND batch_size * max_seqlen % (128 * tp_size) == 0.
-        # Compute the required alignment for max_seqlen:
-        fp8_total_align = 128 * tp_size
-        fp8_seq_align = fp8_total_align // math.gcd(batch_size, fp8_total_align)
-        # Also ensure tp alignment for SP
-        fp8_seq_align = math.lcm(fp8_seq_align, tp_size)
-        max_seqlen = ((max_seqlen + fp8_seq_align - 1) // fp8_seq_align) * fp8_seq_align
+        max_seqlen = _align_bshd_max_seqlen_for_fp8(
+            max_seqlen=max_seqlen,
+            batch_size=batch_size,
+            tp_size=tp_size,
+            cp_size=cp_size,
+        )
 
     attention_mask = torch.zeros(batch_size, max_seqlen, dtype=torch.bool, device=input_ids.device)
     input_ids_bshd = torch.zeros(batch_size, max_seqlen, dtype=input_ids.dtype, device=input_ids.device)