refactor(packing): replace global HF monkeypatch with typed metadata

[beccohov]

What does this PR do ?

Replaces the process-wide monkeypatching that NEAT packing used to route packed sequences through FlashAttention with typed metadata that the collaters build and pass through the public HF kwargs.

Changelog

• Add PackedSeqParams and packed_seq_params_from_doc_ids (datasets/packed_seq.py). The collater derives cu_seqlens/max_seqlen from the document map and emits them as FlashAttentionKwargs (cu_seq_lens_q/k, max_length_q/k). The cu_seqlens span the full flattened batch, including padding, with a forced boundary at every row so a row-filling document cannot merge with the next row.
• LLM and VLM NEAT collaters stop emitting an indexed attention mask on the flash path. Without a mask HF takes its varlen-kwargs branch instead of unpadding a binary one. The per-document map is still exposed as _packed_seq_ids for the loss and context-parallel consumers.
• Remove configure_packing and the three patch sites (_get_unpad_data, _preprocess_mask_arguments, per-module create_causal_mask). In their place validate_flash_packing_support fails early when the installed Transformers build or the model cannot accept the typed contract.
• SDPA/eager keep the existing 4D block-causal path unchanged.
• Tests cover the metadata, both collater backends, the validation failure modes, and forward-pass parity showing no cross-document leakage: block-causal on CPU, plus a flash-varlen check on GPU for padded and unpadded packs.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? No user-facing behavior or config changed.

Additional Information

• Related to refactor(packing): replace global HF monkeypatch with typed metadata #3107

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[beccohov]

Hey, @yuhezhang-ai
as promised, here is the PR. I believe I understood your idea correctly!
I run tests locally and it works correctly.

[akoumpa]

/ok to test d81ddcf

[yuhezhang-ai]

Thanks for taking this on — replacing the process-wide monkeypatches with explicit packed-sequence metadata is a much cleaner direction. I had two follow-ups that would make this safer to rely on broadly:

1. In _validate_model_consumes_packed_contract, could we require either **kwargs, _packed_seq_ids, or all four varlen kwargs? The current set intersection accepts a forward that exposes only one of the four. filter_forward_kwargs can then drop the remaining cumulative-length fields, while Transformers enters its varlen path only when all four are present. Using set(_PACKED_VARLEN_KWARGS).issubset(params) for that branch would keep this check fail-closed.

2. Could we narrow the FlashAttention availability probe so unexpected model/CUDA execution failures are not converted into skips? The real GPU parity test currently covers Llama with a single packed row. It would also be valuable to exercise batch size greater than one with uneven row padding and the actual Qwen3-VL packed FlashAttention path. The release matrix should auto-discover the Qwen3-VL config; we would also be happy to help run the relevant scoped CI job after the PR is updated.

The focused CPU tests passed locally; the FlashAttention cases skipped in the CPU-only environment. Thanks again!

[beccohov]

Thanks, @yuhezhang-ai, I believe I've addressed all of your suggestions:

1. Contract check -- a forward that exposes only some of the four no longer slips through.
2. FA availability. Instead of narrowing the probe, I dropped the forward-based probe entirely and gate on importlib.util.find_spec("flash_attn") + CUDA, matching test_eagle3_packing_fa2_parity.py. There's no probe forward left, so nothing can turn a real model/CUDA failure into a skip.
3. Coverage. Added batch > 1 with uneven per-row padding, and a Qwen3-VL packed FlashAttention parity check that asserts per-document logits match the isolated forward. Left the full-size Qwen3-VL config to the release matrix, as you offered.

I think we need to run CI

[akoumpa]

/ok to test 9bde8b1