Pass packed-document boundaries to HF attention so FSDP runs on FlashAttention-4 by HJSang · Pull Request #4316 · radixark/miles

HJSang · 2026-10-22T21:00:13Z

Stack (merge in order): #4308 → #4315 → #4316 → #4311 (needs radixark/miles-wheels#4).

Problem

Only the GatedDeltaNet adapters passed the model the packed-document boundaries get_batch records. For every other arch, HF re-derived them from position_ids in each attention layer, with max_length as a 0-d tensor. FlashAttention-2 accepts that; FlashAttention-4 raises Expected IntTuple, whose leaves are integers, but got 1223. FA4 is the only training-capable attention a pip install has without a flash-attn build: SGLang depends on it, and its own FA3 (in sglang-kernel) is forward-only.

Change

hf_attention_packing_kwargs: cu_seq_lens_q/k, max_length_q/k (an int).
ArchAdapter.passes_attention_packing_kwargs, set on Qwen3, Qwen3-MoE, Qwen3-VL and GLM-4.7-Flash (stock HF modeling), sends them under every kernel. Remote code and archs without an adapter keep {}; NemotronH is unchanged.
The boundaries are exactly the ones HF derives from position_ids: the documents, then each pad token as its own 1-token segment, built on the device from host lengths (no sync). max_length is then the longest real document and covers every segment. get_batch keeps the pad tail as one segment; sending that with its pad-inflated max_seqlen changed FlashAttention's launch plan and broke true-on-policy's bitwise parity under FA3 (log_probs -1.38827 vs rollout -1.38723), and a real-document max with the one-segment tail would leave pad rows unwritten. Qwen3.5 keeps get_batch's boundaries and max_seqlen as on main.
The vision-tower guard moves to packing.py; Qwen3-VL applies it like Qwen3.5.
Docs: --attn-implementation flash_attention_4, and which architectures can use it.

With FlashAttention-2/3 these archs now take HF's varlen-kwargs branch instead of re-deriving the boundaries per layer, with the same boundaries for real tokens.

Testing

On an H200 devbox with the miles CI image, at this PR's head:

tests/e2e/fsdp/test_qwen3_0_6b_fsdp_sampling_support_replay_true_on_policy.py (FlashAttention-3, boundaries sent) passes: train_rollout_logprob_abs_diff 0.0 on all 5 rollouts.
test_fsdp_hf_compat.py, test_fsdp_hub_kernels.py, fsdp_utils/test_actor.py, test_fsdp_attn_bridge.py and tests/fast-gpu/test_fsdp_gated_deltanet_packing.py: 94 pass. Three CPU-tensor Qwen3.5 tests fail identically on unmodified main in that image (its Triton / causal_conv1d kernels reject CPU tensors); the hosted CPU stage passes them.
New tests: each adapter's kwargs; every Qwen3 attention layer receives the boundaries with an int max_length; attention boundaries equal transformers' own prepare_fa_kwargs_from_position_ids (pad 0, 1, 7); the actor gives stock HF attention each pad token as its own segment, without a device sync (GPU); Qwen3-VL's tower never sees them.
End to end, wheels only: FSDP GRPO on Qwen3-0.6B with flash_attention_4 (exit 0, log-prob gap 0.0199 → 0.0157, grad norm 0.36 / 0.43), matching a FlashAttention-2 run. That run predates the real-max change.
Needs run-ci-fsdp before merge (label on).

🤖 Generated with Claude Code

HJSang · 2026-10-22T21:00:26Z

@claude review always

HJSang · 2026-10-22T21:00:28Z

@codex review

chatgpt-codex-connector · 2026-10-22T21:00:46Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-22T21:04:34.957264Z	`9a32f70`	Manual request

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

chatgpt-codex-connector · 2026-10-22T21:04:32Z

Codex Review: Didn't find any major issues. Swish!

Reviewed commit: 9a32f708db

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

claude

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

HJSang · 2026-10-22T21:05:22Z

@claude review

claude

Code review found no issues

No high-confidence issues detected in this change.

claude

Code review found no issues

No high-confidence issues detected in this change.

claude

Code review found no issues

No high-confidence issues detected in this change.

claude

Code review found no issues

No high-confidence issues detected in this change.

claude

Code review found no new issues

No new issues were found in this update; 1 finding from earlier reviews is still open above.

Still open from earlier reviews (1):

🔴 miles/backends/fsdp_utils/actor.py:410 — Packed micro-batches whose pad tail is longer than every real document can now get uninitialized attention outputs or g…

If you have decided not to act on one of these findings, resolve its thread (a reply alone leaves it open) and the next review stops counting it. To review this commit again now, use Re-run on its "Claude Code Review" check.

…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>

claude

Code review found no issues

No high-confidence issues detected in this change.

…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>

Only the GatedDeltaNet adapters handed the model the boundaries get_batch records; every other arch left HF to re-derive them from position_ids in each attention layer. That fallback gives the kernel max_length as a 0-d tensor, which FlashAttention-2 tolerates and FlashAttention-4 rejects, so FSDP could not train on the FlashAttention-4 that SGLang installs, the only training-capable attention a pip install without a flash-attn build has. Adapters for stock HF modeling (Qwen3, Qwen3-MoE, Qwen3-VL, GLM-4.7-Flash) now opt in with passes_attention_packing_kwargs and send cu_seq_lens_q/k and max_length_q/k, with max_length a Python int. Remote code and unknown archs keep their stock inputs. Qwen3-VL, like Qwen3.5, keeps them out of its vision tower; that guard moves to packing.py beside the kwargs it drops. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

transformers resolves any *flash* attn_implementation to a FlashAttention kernel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

With them, Qwen3 on FlashAttention-3 took HF's varlen-kwargs branch instead of re-deriving the boundaries from position_ids, and the true-on-policy e2e test lost bitwise parity with SGLang (log_probs -1.38827 vs rollout -1.38723). Only FlashAttention-4 rejects the re-derived boundaries, so only it gets them now; every other kernel keeps HF's path as on main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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>

…ntion-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>

…on-4 needs The actor test's pad segment was as long as its longest document, so it could not tell get_batch's max_seqlen from the real one. The docs now list every architecture that passes the boundaries and say the rest cannot use flash_attention_4 yet. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…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>

claude

Code review found no issues

No high-confidence issues detected in this change.

guapisolo

cc cmt

Approving #4316 at cecb2b4. No blocking findings in this increment. The review ran at e3fa3d7; this head is the same patch rebased, plus main's config-access-exempt comments on the moved keep_packing_kwargs_out_of_vision.

Verification: read-only review of this increment at e3fa3d7 with its green PR Test run (40553236168), and a comparison of the incremental diffs at both heads. Not covered: the checks still running at this head, a FlashAttention-4 end-to-end run with the final boundary logic, and bitwise parity of a single unpadded document between the dense and varlen FA2/FA3 paths.

HJSang added the run-ci-fsdp label Oct 22, 2026

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from 9a32f70 to 25ac4f4 Compare October 22, 2026 21:04

HJSang changed the base branch from main to hjsang/native-vocab-ce October 22, 2026 21:04

This was referenced Oct 22, 2026

Build miles as a pure wheel that imports FSDP and torchtitan without Megatron #4308

Open

Install miles-train and miles-train-async with the wheel #4312

Closed

Compute log probs with one plain-torch vocab-parallel cross entropy, without Megatron #4315

Open

HJSang marked this pull request as ready for review October 22, 2026 21:05

HJSang requested review from Rockdu, Zhichenzzz, fzyzcjy, guapisolo and yueming-yuan as code owners October 22, 2026 21:05

claude Bot reviewed Oct 22, 2026

View reviewed changes

HJSang mentioned this pull request Oct 22, 2026

Build the sglang-miles wheel hourly and keep every version #4311

Open

HJSang force-pushed the hjsang/native-vocab-ce branch from 1fa6b28 to df4214c Compare October 22, 2026 21:08

HJSang requested review from Shi-Dong and maocheng23 as code owners October 22, 2026 21:08

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from 25ac4f4 to ac61006 Compare October 22, 2026 21:08

claude Bot reviewed Oct 22, 2026

View reviewed changes

HJSang force-pushed the hjsang/native-vocab-ce branch from df4214c to f8cd0e7 Compare October 22, 2026 21:20

HJSang requested review from jybsuper and yushengsu-thu as code owners October 22, 2026 21:20

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from ac61006 to 9ec3639 Compare October 22, 2026 21:20

claude Bot reviewed Oct 22, 2026

View reviewed changes

HJSang force-pushed the hjsang/native-vocab-ce branch from f8cd0e7 to 2bf5e05 Compare October 22, 2026 22:17

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from 9ec3639 to 6b216e2 Compare October 22, 2026 22:17

HJSang force-pushed the hjsang/native-vocab-ce branch from 2bf5e05 to 2b84a30 Compare October 22, 2026 22:21

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from 6b216e2 to 1dac250 Compare October 22, 2026 22:21

HJSang changed the title ~~Pass packed-document boundaries to stock HF attention in FSDP~~ Pass packed-document boundaries to HF attention so FSDP runs on FlashAttention-4 Oct 22, 2026

claude Bot reviewed Oct 22, 2026

View reviewed changes

claude Bot reviewed Oct 26, 2026

View reviewed changes

Comment thread miles/backends/fsdp_utils/actor.py Outdated

claude Bot reviewed Oct 26, 2026

View reviewed changes

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from 52990f6 to e3fa3d7 Compare October 27, 2026 17:43

HJSang and others added 8 commits October 27, 2026 14:23

test: give the Qwen3 attention stand-in a non-flash name

1e28716

transformers resolves any *flash* attn_implementation to a FlashAttention kernel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Trim the comments and docs in the FSDP packing-kwargs change

82828d3

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

HJSang force-pushed the hjsang/native-vocab-ce branch from a53d9a6 to 9022cb2 Compare October 27, 2026 21:34

HJSang force-pushed the hjsang/fsdp-default-packing-kwargs branch from e3fa3d7 to cecb2b4 Compare October 27, 2026 21:34

claude Bot reviewed Oct 27, 2026

View reviewed changes

guapisolo approved these changes Oct 27, 2026

View reviewed changes

yueming-yuan approved these changes Oct 28, 2026

View reviewed changes

Conversation

HJSang commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Change

Testing

Uh oh!

HJSang commented Oct 22, 2026

Uh oh!

HJSang commented Oct 22, 2026

Uh oh!

chatgpt-codex-connector Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

chatgpt-codex-connector Bot commented Oct 22, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

HJSang commented Oct 22, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

guapisolo left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

HJSang commented Oct 22, 2026 •

edited

Loading

chatgpt-codex-connector Bot commented Oct 22, 2026 •

edited

Loading