Give stock HF attention the boundaries HF derives: each pad token on its own

real_max_seqlen dropped the pad segment from max_length while cu_seqlens still
carried it as one segment of up to 127 tokens; FlashAttention leaves rows past
max_length unwritten, so pad rows could get uninitialized outputs and
gradients. hf_attention_packing_kwargs now passes the documents, then each
pad token as a 1-token segment, exactly what HF's position_ids path derives,
so the longest real document covers every segment. It is built on the device
from host lengths, without a sync. Qwen3.5 keeps get_batch's boundaries and
max_seqlen as on main.

Reported by claude[bot] on #4316.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: HJSang

miles/backends/fsdp_utils/actor.py

@@ -22,7 +22,6 @@
 from miles.utils.workers.rpc.common.wire_types import Pickled
 
 from .adaptations.config_checks import validate_hf_config
-from .adaptations.packing import real_max_seqlen
 from .adaptations.precision import apply_fp32_master, default_precision_policy, precision_forward_context
 from .adaptations.specs import resolve_arch_adapter
 from .hf_weight_iterator import FSDPHfWeightIterator
@@ -407,7 +406,8 @@ def _get_model_inputs_args(self, batch: dict) -> dict:
             self.arch_adapter.packing_kwargs(
                 cu_seqlens=batch["cu_seqlens"],
                 cu_seqlens_host=batch["cu_seqlens_host"],
-                max_seqlen=real_max_seqlen(batch["cu_seqlens_host"], len(batch["unconcat_tokens"])),
+                max_seqlen=batch["max_seqlen"],
+                num_documents=len(batch["unconcat_tokens"]),
             )
         )
 

miles/backends/fsdp_utils/adaptations/arch_adapter.py

@@ -41,13 +41,15 @@ def patch_model(self, model: nn.Module, args: Namespace) -> None:
         and the ref model, so it must be idempotent."""
 
     def packing_kwargs(
-        self, *, cu_seqlens: torch.Tensor, cu_seqlens_host: tuple[int, ...], max_seqlen: int
+        self, *, cu_seqlens: torch.Tensor, cu_seqlens_host: tuple[int, ...], max_seqlen: int, num_documents: int
     ) -> dict[str, object]:
         """Extra forward kwargs that tell this architecture's modeling where the packed documents that
         ``get_batch`` recorded start."""
         if not self.passes_attention_packing_kwargs:
             return {}
-        return hf_attention_packing_kwargs(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+        return hf_attention_packing_kwargs(
+            cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, num_documents=num_documents
+        )
 
     def param_transform(self, name: str, param: torch.Tensor) -> ParamExpand | None:
         """How to rewrite ``param`` for the rollout engine at weight sync, or None to stream it unchanged."""

miles/backends/fsdp_utils/adaptations/packing.py

@@ -41,20 +41,25 @@ def packed_seq_context(position_ids: torch.Tensor | None) -> PackedSeqContext |
 HF_PACKING_KWARG_NAMES = frozenset({"cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k", "seq_idx"})
 
 
-def real_max_seqlen(cu_seqlens_host: tuple[int, ...], num_documents: int) -> int:
-    """The longest real document; ``get_batch``'s ``max_seqlen`` also counts the trailing pad segment, which
-    HF never does (each pad token is its own document there) and which changes the kernel's launch plan."""
-    real = cu_seqlens_host[: num_documents + 1]
-    return max(end - start for start, end in itertools.pairwise(real))
-
+def hf_attention_packing_kwargs(
+    *, cu_seqlens: torch.Tensor, cu_seqlens_host: tuple[int, ...], num_documents: int
+) -> dict[str, object]:
+    """The boundaries HF derives from ``position_ids``: the documents, then each pad token on its own.
 
-def hf_attention_packing_kwargs(*, cu_seqlens: torch.Tensor, max_seqlen: int) -> dict[str, object]:
-    """The varlen boundaries HF attention reads; ``max_length`` stays an int, which FlashAttention-4 needs."""
+    ``max_length`` is an int, since FlashAttention-4 rejects HF's tensor, and it covers every segment, since
+    the kernel leaves rows past it unwritten. Keeping the pad tail as one segment would raise it, and the
+    kernel's launch plan with it, away from what HF (and true-on-policy's bitwise match) uses.
+    """
+    pad = cu_seqlens_host[-1] - cu_seqlens_host[num_documents]
+    # on the device from host lengths, so building it never syncs
+    pad_ends = cu_seqlens[num_documents] + torch.arange(1, pad + 1, dtype=cu_seqlens.dtype, device=cu_seqlens.device)
+    boundaries = torch.cat([cu_seqlens[: num_documents + 1], pad_ends])
+    max_length = max(end - start for start, end in itertools.pairwise(cu_seqlens_host[: num_documents + 1]))
     return {
-        "cu_seq_lens_q": cu_seqlens,
-        "cu_seq_lens_k": cu_seqlens,
-        "max_length_q": max_seqlen,
-        "max_length_k": max_seqlen,
+        "cu_seq_lens_q": boundaries,
+        "cu_seq_lens_k": boundaries,
+        "max_length_q": max_length,
+        "max_length_k": max_length,
     }
 
 
@@ -73,7 +78,10 @@ def hf_packing_kwargs(
         output_size=cu_seqlens_host[-1],
     )
     return {
-        **hf_attention_packing_kwargs(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen),
+        "cu_seq_lens_q": cu_seqlens,
+        "cu_seq_lens_k": cu_seqlens,
+        "max_length_q": max_seqlen,
+        "max_length_k": max_seqlen,
         "seq_idx": seq_idx.unsqueeze(0),
     }
 

miles/backends/fsdp_utils/adaptations/specs/qwen3_5.py

@@ -10,7 +10,7 @@
 class Qwen35Adapter(ArchAdapter):
     model_types = frozenset({"qwen3_5", "qwen3_5_text", "qwen3_next"})
 
-    def packing_kwargs(self, *, cu_seqlens, cu_seqlens_host, max_seqlen):
+    def packing_kwargs(self, *, cu_seqlens, cu_seqlens_host, max_seqlen, num_documents):
         # GatedDeltaNet reads seq_idx (causal conv) and cu_seq_lens_q (FLA chunk rule) from **kwargs.
         return hf_packing_kwargs(cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen)
 

tests/fast-gpu/test_fsdp_gated_deltanet_packing.py

@@ -298,6 +298,24 @@ def test_building_packing_kwargs_never_syncs_the_device():
     assert inputs["cu_seq_lens_q"] is batch["cu_seqlens"] and inputs["max_length_q"] == 129
 
 
+def test_building_attention_boundaries_never_syncs_the_device():
+    """Stock HF attention gets the pad tail as 1-token segments, built on the device from host lengths."""
+    actor = object.__new__(FSDPTrainRayActor)
+    actor.arch_adapter = resolve_arch_adapter(SimpleNamespace(model_type="qwen3"))
+    docs = [_random_tokens(length, seed) for seed, length in enumerate((37, 64, 129))]
+    batch = _packed_batch(docs, padded_length=256)
+
+    previous_mode = torch.cuda.get_sync_debug_mode()
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        inputs = actor._get_model_inputs_args(batch)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous_mode)
+
+    assert inputs["cu_seq_lens_q"].tolist() == [0, 37, 101, 230, *range(231, 257)]
+    assert inputs["max_length_q"] == 129
+
+
 if __name__ == "__main__":
     import sys
 

tests/fast/backends/fsdp_utils/test_actor.py

@@ -91,22 +91,40 @@ def test_packing_kwargs_follow_the_batch_boundaries_so_padding_stays_one_segment
         """Pad tokens carry position 0; boundaries derived from `position_ids` would make each one a document."""
         actor = object.__new__(actor_module.FSDPTrainRayActor)
         actor.arch_adapter = resolve_arch_adapter(SimpleNamespace(model_type="qwen3_5_text"))
-        cu_seqlens_host = (0, 3, 5, 9)  # two documents, then the 4-token pad segment get_batch appends
+        cu_seqlens_host = (0, 3, 5, 8)  # two documents, then the 3-token pad segment get_batch appends
+        batch = {
+            "tokens": torch.zeros(1, 8, dtype=torch.long),
+            "position_ids": torch.tensor([[0, 1, 2, 0, 1, 0, 0, 0]]),
+            "cu_seqlens": torch.tensor(cu_seqlens_host, dtype=torch.int32),
+            "cu_seqlens_host": cu_seqlens_host,
+            "max_seqlen": 3,
+            "unconcat_tokens": [torch.zeros(3, dtype=torch.long), torch.zeros(2, dtype=torch.long)],
+        }
+
+        model_args = actor._get_model_inputs_args(batch)
+
+        assert model_args["cu_seq_lens_q"] is batch["cu_seqlens"]
+        assert model_args["max_length_q"] == 3
+        assert model_args["seq_idx"].tolist() == [[0, 0, 0, 1, 1, 2, 2, 2]]
+
+    def test_stock_hf_attention_gets_every_pad_token_as_its_own_segment(self) -> None:
+        """A 4-token pad tail kept as one segment would exceed max_length 3, leaving its rows unwritten."""
+        actor = object.__new__(actor_module.FSDPTrainRayActor)
+        actor.arch_adapter = resolve_arch_adapter(SimpleNamespace(model_type="qwen3"))
+        cu_seqlens_host = (0, 3, 5, 9)
         batch = {
             "tokens": torch.zeros(1, 9, dtype=torch.long),
             "position_ids": torch.tensor([[0, 1, 2, 0, 1, 0, 0, 0, 0]]),
             "cu_seqlens": torch.tensor(cu_seqlens_host, dtype=torch.int32),
             "cu_seqlens_host": cu_seqlens_host,
-            "max_seqlen": 4,  # get_batch's, which counts the pad segment
+            "max_seqlen": 4,
             "unconcat_tokens": [torch.zeros(3, dtype=torch.long), torch.zeros(2, dtype=torch.long)],
         }
 
         model_args = actor._get_model_inputs_args(batch)
 
-        assert model_args["cu_seq_lens_q"] is batch["cu_seqlens"]
-        # the longest real document: the pad segment changes FlashAttention's launch plan
+        assert model_args["cu_seq_lens_q"].tolist() == [0, 3, 5, 6, 7, 8, 9]
         assert model_args["max_length_q"] == 3
-        assert model_args["seq_idx"].tolist() == [[0, 0, 0, 1, 1, 2, 2, 2, 2]]
 
 
 class TestFSDPTrainParallelConfigWiring:

tests/fast/backends/test_fsdp_hf_compat.py

@@ -26,13 +26,14 @@ def _adapter(model_type):
     return resolve_arch_adapter(SimpleNamespace(model_type=model_type))
 
 
-def _packed_boundaries(*lengths):
-    """The boundary fields `get_batch` records for documents of these lengths."""
-    cu_seqlens_host = tuple(itertools.accumulate(lengths, initial=0))
+def _packed_boundaries(*lengths, pad=0):
+    """The boundary fields `get_batch` records for documents of these lengths, then its pad segment."""
+    cu_seqlens_host = tuple(itertools.accumulate([*lengths, *([pad] if pad else [])], initial=0))
     return dict(
         cu_seqlens=torch.tensor(cu_seqlens_host, dtype=torch.int32),
         cu_seqlens_host=cu_seqlens_host,
-        max_seqlen=max(lengths),
+        max_seqlen=max(lengths + (pad,)),
+        num_documents=len(lengths),
     )
 
 
@@ -338,7 +339,9 @@ def test_hf_packing_kwargs_match_the_padding_free_collator():
 
     collator = DataCollatorWithFlattening(return_tensors="pt", return_flash_attn_kwargs=True, return_seq_idx=True)
     flattened = collator([{"input_ids": [1, 2, 3]}, {"input_ids": [4, 5]}, {"input_ids": [6, 7, 8, 9]}])
-    kwargs = hf_packing_kwargs(**_packed_boundaries(3, 2, 4))
+    boundaries = _packed_boundaries(3, 2, 4)
+    del boundaries["num_documents"]
+    kwargs = hf_packing_kwargs(**boundaries)
 
     assert set(kwargs) == HF_PACKING_KWARG_NAMES
     for name, value in kwargs.items():
@@ -349,12 +352,22 @@ def test_hf_packing_kwargs_match_the_padding_free_collator():
             assert value == expected, name
 
 
-def test_real_max_seqlen_leaves_out_the_pad_segment():
-    """A 7-token pad tail must not become the max: HF never counts it, and the kernel plan follows it."""
-    from miles.backends.fsdp_utils.adaptations.packing import real_max_seqlen
+@pytest.mark.parametrize("pad", [0, 1, 7])
+def test_attention_boundaries_are_the_ones_hf_derives_from_position_ids(pad):
+    """True-on-policy's bitwise match relies on HF's own boundaries; a max_length below a segment's length
+    would leave its rows unwritten, as a 7-token pad tail kept as one segment would."""
+    from transformers.modeling_flash_attention_utils import prepare_fa_kwargs_from_position_ids
 
-    assert real_max_seqlen((0, 3, 5, 12), num_documents=2) == 3
-    assert real_max_seqlen((0, 3, 5), num_documents=2) == 3  # no pad segment
+    from miles.backends.fsdp_utils.adaptations.packing import hf_attention_packing_kwargs
+
+    boundaries = _packed_boundaries(3, 2, pad=pad)
+    del boundaries["max_seqlen"]
+    kwargs = hf_attention_packing_kwargs(**boundaries)
+    position_ids = torch.tensor([[0, 1, 2, 0, 1, *([0] * pad)]])
+    (hf_cu_seqlens, _), (hf_max_length, _) = prepare_fa_kwargs_from_position_ids(position_ids)
+
+    assert kwargs["cu_seq_lens_q"].tolist() == hf_cu_seqlens.tolist()
+    assert type(kwargs["max_length_q"]) is int and kwargs["max_length_q"] == int(hf_max_length) == 3
 
 
 def test_each_arch_passes_the_packing_kwargs_its_modeling_reads():