Give attention the longest real document as max_length

get_batch's max_seqlen also counts the trailing pad segment, which HF's own
position_ids path never does, and FlashAttention picks its launch plan from
max_seqlen. batch["max_seqlen"] itself is unchanged (Megatron reads it).

Found by an adversarial review of this PR.

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

Author: HJSang

miles/backends/fsdp_utils/actor.py

@@ -22,6 +22,7 @@
 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
@@ -406,7 +407,7 @@ 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=batch["max_seqlen"],
+                max_seqlen=real_max_seqlen(batch["cu_seqlens_host"], len(batch["unconcat_tokens"])),
                 attn_implementation=self.args.attn_implementation,
             )
         )

miles/backends/fsdp_utils/adaptations/packing.py

@@ -8,6 +8,7 @@
 """
 
 import functools
+import itertools
 from dataclasses import dataclass
 
 import torch
@@ -40,6 +41,13 @@ 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, max_seqlen: int) -> dict[str, object]:
     """The varlen boundaries HF attention reads; ``max_length`` stays an int, which FlashAttention-4 needs."""
     return {

tests/fast-gpu/test_fsdp_gated_deltanet_packing.py

@@ -116,6 +116,7 @@ def _packed_batch(docs, padded_length):
         cu_seqlens=torch.tensor(cu_seqlens_host, dtype=torch.int32, device="cuda"),
         cu_seqlens_host=cu_seqlens_host,
         max_seqlen=max(end - start for start, end in itertools.pairwise(cu_seqlens_host)),
+        unconcat_tokens=docs,
     )
 
 

tests/fast/backends/fsdp_utils/test_actor.py

@@ -99,6 +99,7 @@ def test_packing_kwargs_follow_the_batch_boundaries_so_padding_stays_one_segment
             "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)

tests/fast/backends/test_fsdp_hf_compat.py

@@ -349,6 +349,14 @@ 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
+
+    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
+
+
 def test_each_arch_passes_the_packing_kwargs_its_modeling_reads():
     boundaries = _packed_boundaries(2, 3)
     for model_type in ("qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", "qwen3_next"):