Send the packed boundaries for every kernel again, not only FlashAttention-4

This reverts the FlashAttention-4-only gate. True-on-policy's FA3 bitwise
break came from max_length counting the pad segment, not from the boundaries
themselves: with the longest real document as max_length (previous commit)
and the boundaries sent under every kernel, the true-on-policy e2e test is
bit-identical again (H200, train_rollout_logprob_abs_diff 0.0 on all five
rollouts). Without the gate, the kernels-community/flash-attn4 spelling of
--attn-implementation gets the boundaries too.

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

Author: HJSang

docs/user-guide/training-backend.md

@@ -381,8 +381,8 @@ targets are Megatron-only: both `--offload-train-target disk` and
 - `--attn-implementation` picks the `transformers` attention backend: `flash_attention_2` by
   default, with `flash_attention_3`, `sdpa` and `eager` passed straight through.
   `flash_attention_4` uses the FlashAttention-4 that SGLang installs, so it trains without a
-  `flash-attn` build. With it, the adapter passes the packed-document boundaries (Qwen3,
-  Qwen3-MoE, Qwen3-VL, GLM-4.7-Flash and Qwen3.5); other kernels keep `transformers`' own path.
+  `flash-attn` build. It needs an adapter that passes the packed-document boundaries: Qwen3,
+  Qwen3-MoE, Qwen3-VL, GLM-4.7-Flash and Qwen3.5.
 - An architecture with fussier numerics can register its own policy, see
   [when an HF model needs help](#going-deeper-when-an-hf-model-needs-help).
 

miles/backends/fsdp_utils/actor.py

@@ -408,7 +408,6 @@ def _get_model_inputs_args(self, batch: dict) -> dict:
                 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"])),
-                attn_implementation=self.args.attn_implementation,
             )
         )
 

miles/backends/fsdp_utils/adaptations/arch_adapter.py

@@ -41,13 +41,11 @@ 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, attn_implementation: str
+        self, *, cu_seqlens: torch.Tensor, cu_seqlens_host: tuple[int, ...], max_seqlen: int
     ) -> dict[str, object]:
         """Extra forward kwargs that tell this architecture's modeling where the packed documents that
         ``get_batch`` recorded start."""
-        # only FlashAttention-4 rejects the boundaries HF re-derives; other kernels keep that path, which
-        # true-on-policy matches bit for bit
-        if not (self.passes_attention_packing_kwargs and attn_implementation == "flash_attention_4"):
+        if not self.passes_attention_packing_kwargs:
             return {}
         return hf_attention_packing_kwargs(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
 

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, attn_implementation):
+    def packing_kwargs(self, *, cu_seqlens, cu_seqlens_host, max_seqlen):
         # 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

@@ -123,7 +123,6 @@ def _packed_batch(docs, padded_length):
 def _actor_for(model):
     actor = object.__new__(FSDPTrainRayActor)
     actor.arch_adapter = resolve_arch_adapter(model.config)
-    actor.args = SimpleNamespace(attn_implementation=model.config._attn_implementation)
     return actor
 
 
@@ -284,7 +283,6 @@ def test_building_packing_kwargs_never_syncs_the_device():
     """A host-device sync on every micro-batch from sizing `seq_idx` with a device-side length."""
     actor = object.__new__(FSDPTrainRayActor)
     actor.arch_adapter = resolve_arch_adapter(SimpleNamespace(model_type="qwen3_5"))
-    actor.args = SimpleNamespace(attn_implementation="flash_attention_2")
     docs = [_random_tokens(length, seed) for seed, length in enumerate((37, 64, 129))]
     batch = _packed_batch(docs, padded_length=256)
 

tests/fast/backends/fsdp_utils/test_actor.py

@@ -91,7 +91,6 @@ 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"))
-        actor.args = SimpleNamespace(attn_implementation="flash_attention_2")
         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),

tests/fast/backends/test_fsdp_hf_compat.py

@@ -360,20 +360,15 @@ def test_real_max_seqlen_leaves_out_the_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"):
-        kwargs = _adapter(model_type).packing_kwargs(**boundaries, attn_implementation="flash_attention_2")
-        assert "seq_idx" in kwargs, model_type
-    # stock HF attention: varlen boundaries for FlashAttention-4 only, max_length an int; other kernels
-    # keep HF's position_ids path, which true-on-policy matches bit for bit
+        assert "seq_idx" in _adapter(model_type).packing_kwargs(**boundaries), model_type
+    # stock HF attention: varlen boundaries, max_length an int
     for model_type in ("glm4_moe_lite", "qwen3", "qwen3_moe", "qwen3_vl"):
-        adapter = _adapter(model_type)
-        kwargs = adapter.packing_kwargs(**boundaries, attn_implementation="flash_attention_4")
+        kwargs = _adapter(model_type).packing_kwargs(**boundaries)
         assert set(kwargs) == {"cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k"}, model_type
         assert type(kwargs["max_length_q"]) is int and kwargs["max_length_q"] == 3, model_type
-        for other in ("flash_attention_2", "flash_attention_3", "sdpa"):
-            assert adapter.packing_kwargs(**boundaries, attn_implementation=other) == {}, (model_type, other)
     # NemotronH resets through its own patch; other archs (including remote code) keep their stock inputs.
     for model_type in ("nemotron_h", "some_remote_lm"):
-        assert _adapter(model_type).packing_kwargs(**boundaries, attn_implementation="flash_attention_4") == {}
+        assert _adapter(model_type).packing_kwargs(**boundaries) == {}, model_type
 
 
 def test_qwen3_attention_receives_the_packed_boundaries(monkeypatch):
@@ -406,7 +401,7 @@ def attention_stand_in(module, query, key, value, attention_mask, scaling=None,
         model(
             input_ids=torch.arange(5).view(1, 5),
             position_ids=torch.tensor([[0, 1, 2, 0, 1]]),
-            **_adapter("qwen3").packing_kwargs(**_packed_boundaries(3, 2), attn_implementation="flash_attention_4"),
+            **_adapter("qwen3").packing_kwargs(**_packed_boundaries(3, 2)),
         )
 
     assert len(seen) == cfg.num_hidden_layers
@@ -432,8 +427,7 @@ def forward(self, pixels, **kwargs):
     adapter.patch_model(model, None)
     assert model.model.visual.forward is wrapped  # idempotent across the actor and ref model
 
-    packing_kwargs = adapter.packing_kwargs(**_packed_boundaries(3, 2), attn_implementation="flash_attention_4")
-    model.model.visual(None, grid_thw=1, **packing_kwargs)
+    model.model.visual(None, grid_thw=1, **adapter.packing_kwargs(**_packed_boundaries(3, 2)))
     assert calls == [{"grid_thw": 1}]
 
 
@@ -480,7 +474,7 @@ def chunk_rule(*args, cu_seqlens=None, **kwargs):
     monkeypatch.setattr(modeling_qwen3_5, "causal_conv1d_fn", conv)
     monkeypatch.setattr(modeling_qwen3_5, "torch_chunk_gated_delta_rule", chunk_rule)
 
-    packing_kwargs = _adapter("qwen3_5_text").packing_kwargs(**_packed_boundaries(3, 2), attn_implementation="eager")
+    packing_kwargs = _adapter("qwen3_5_text").packing_kwargs(**_packed_boundaries(3, 2))
     with torch.no_grad():
         model(input_ids=torch.arange(5).view(1, 5), position_ids=torch.tensor([[0, 1, 2, 0, 1]]), **packing_kwargs)
 
@@ -533,7 +527,7 @@ def flash_stand_in(module, query, key, value, attention_mask, scaling=None, cu_s
         position_ids=torch.tensor([[0, 1, 2, 0, 1, 2, 3]]),
         pixel_values=torch.randn(4, 3 * 2 * 2 * 2),
         image_grid_thw=torch.tensor([[1, 2, 2]]),
-        **adapter.packing_kwargs(**_packed_boundaries(3, 4), attn_implementation="eager"),
+        **adapter.packing_kwargs(**_packed_boundaries(3, 4)),
     )
     with torch.no_grad():
         model(**inputs)

tests/fast/backends/test_fsdp_hub_kernels.py

@@ -524,7 +524,7 @@ def hub_conv(
     inputs = dict(
         input_ids=torch.arange(5).view(1, 5),
         position_ids=torch.tensor([[0, 1, 2, 0, 1]]),
-        **_adapter("qwen3_5_text").packing_kwargs(**_packed_boundaries(3, 2), attn_implementation="eager"),
+        **_adapter("qwen3_5_text").packing_kwargs(**_packed_boundaries(3, 2)),
     )
     with torch.no_grad():
         native_logits = model(**inputs).logits