[chsu-amd]

What does this PR do?

Fixes #15251

Related: #13375 (pipeline DtoH syncs), #13401 (pipeline profiling), #13634 (Lumina review, different scope; see self-review notes). Lumina-Next-SFT (LuminaPipeline + LuminaNextDiT2DModel) repeats avoidable work on every step. This PR removes it in two commits:

1. LuminaAttnProcessor2_0: stop expanding the (batch, 1, 1, kv_len) bool mask to (batch, heads, q_len, kv_len). F.scaled_dot_product_attention broadcasts the mask.
2. LuminaPatchEmbed: slice freqs_cis to the token grid before .to(device). LuminaPipeline: build the linear-scaled and NTK-scaled RoPE tables once per call on the latent token grid and move them to the device before the denoising loop, instead of building a 384 x 384 table on CPU every step and moving all of it. The loop only picks one of the two tables at scaling_watershed, as before.

No new branches or flags. Output is bit-identical (0 differing values) at 1024x1024 (card setting) and at 512x512 and 768x1024 with scaling_watershed=0.3, where the loop switches from the linear-scaled to the NTK-scaled table.

LuminaFeedForward (shared with Lumina-Image-2.0) is changed separately (#15252).

Full pipeline call, bf16, seconds, main vs this PR:

	main	this PR	speedup
RTX 5090, eager	11.928	10.003	1.192x
RTX 5090, transformer.compile(fullgraph=True)	10.300	8.784	1.173x
AMD Radeon AI PRO R9700 (gfx1201), eager	26.420	16.059	1.645x

Alpha-VLLM/Lumina-Next-SFT-diffusers, model card prompt and settings (from_pretrained(torch_dtype=torch.bfloat16).to("cuda"), __call__ defaults: 1024x1024, 30 steps, guidance 4.0), CPU generator with a fixed seed. RTX 5090 (Windows), bf16, batch 1, default attention backend (SDPA via LuminaAttnProcessor2_0), torch 2.14.1+cu130. AMD Radeon AI PRO R9700 (gfx1201): Linux, torch 2.16.0a0+rocm10.3, same settings.

Method: installed diffusers 0.42.1, whose touched classes (LuminaAttnProcessor2_0, LuminaPatchEmbed, LuminaPipeline) equal main 0905b66, with main's and this PR's functions alternated in one process. Each row is one process: 8 pairs (main first in every pair), median per side (mean of the 4th and 5th of 8 sorted times); the change was faster in 8 of 8 pairs in every row. 0 differing output values vs main for seeds 0, 1, 7 and in every pair. 0 graph breaks under compile.

A plain timing script (seed 0, one process per checkout) is in #15251: 11.968 -> 10.001 s at the card setting, 0 differing values.

cudaStreamSynchronize calls per pipeline call on RTX 5090 (eager, torch.profiler) with commit 2: 74 -> 46.

Tests

Each commit adds the test for its own change in tests/models/transformers/test_models_transformer_lumina.py (TestLuminaNextDiTBehavior, padded prompt mask, non-zero gate so the cross-attention and its mask are not hidden):

• Commit 1: test_attention_mask_broadcast_matches_expanded_mask checks the transformer output with the broadcast mask against SDPA called with the mask expanded to (batch, heads, q_len, kv_len), bit for bit. It also checks that padded text tokens do not change the output, which fails if the mask is dropped from the SDPA call.
• Commit 2: test_rotary_table_on_token_grid_matches_384_table checks that the table built on the token grid equals the [:tokens, :tokens] slice of the 384 x 384 table, also on a non-square (tokens - 2, tokens) grid, and that the transformer output is the same with either, for linear and NTK factors 2.0 and 0.5 and the unscaled table. Its table asserts fail if get_2d_rotary_pos_embed_lumina returns a grid table that is not the slice (checked by shifting it one row and one column, and by swapping height and width on the non-square grid).

Both new tests check equivalence, so they should also pass on main (not run on main). The new pipeline code runs in the fast TestLuminaPipeline tests (32x32, where both tables are equal). The pipeline's own height / width order and linear / NTK pick are covered only by the end-to-end runs with scaling_watershed=0.3 above (768x1024 is the non-square one).

From the checkout root, with that checkout's src first on PYTHONPATH:

python -m pytest -q -p no:cacheprovider tests/models/transformers/test_models_transformer_lumina.py tests/pipelines/lumina/test_lumina_nextdit.py

main,    RTX 5090 (Windows):           53 passed, 10 skipped, 9 warnings in 89.25s
this PR, RTX 5090 (Windows):           55 passed, 10 skipped, 9 warnings in 90.35s
main,    CPU (CUDA_VISIBLE_DEVICES=-1): 33 passed, 30 skipped, 8 warnings in 52.42s
this PR, CPU (CUDA_VISIBLE_DEVICES=-1): 35 passed, 30 skipped, 8 warnings in 49.90s

No failures on either side. The commit-1 tree alone also passes TestLuminaNextDiTBehavior on CUDA and CPU.

Self-review notes (self-review skill)

READY, no blocking issues. ruff check and ruff format --check (ruff 0.9.10) pass on the changed files. utils/check_copies.py and utils/check_dummies.py could not be run cleanly on my Windows host: they fail on unmodified main too, with the same output as on this branch (103 copy inconsistencies in unrelated files such as loaders/lora_pipeline.py, and dummy_bitsandbytes_objects), none in the changed files; no # Copied from block covers the changed code. No dead code. Left as is:

• The mask test patches F.scaled_dot_product_attention to compute the reference with the expanded mask; the real SDPA still runs.
• Both RoPE tables are built every call, even when the loop only uses one of them. That is two small CPU tables per call instead of one per step, and the loop stays a plain pick.
• The loop still syncs once per step on current_timestep[0] < scaling_watershed (device tensor vs float, already on main); not changed here.
• Not touched (already on main): LuminaAttnProcessor2_0 calls SDPA directly instead of dispatch_attention_fn (lumina model/pipeline review #13634, Issue 6), and self-attention passes an all-ones mask instead of None.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or GitHub Discussions? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings. (no user-facing change)
• Did you write any new necessary tests?
• Are you the author (or part of the team we work with) of the model/pipeline this PR relates to?

Who can review?

cc @yiyixuxu @dg845