[Diffusion] Optimize segmented varlen attention packing

[BBuf]

Motivation

Qwen-Image joint attention keeps text and image Q/K/V in separate tensors, then materializes three [text, image] concatenations before the masked varlen path immediately gathers the valid rows again. On a 1024x1024 Qwen-Image-2512 request this adds three dense copies in every one of the 60 transformer blocks.

The same profile also showed that batch-1 packed attention used FA3's VarlenDynamicPersistent scheduler while vLLM-Omni used the faster StaticPersistent scheduler for the equivalent single sequence.

Modifications

• Add a bitwise segmented Triton pack that reads a virtual [prefix, main] Q/K/V sequence without materializing the three concatenations.
• Let USPAttention accept an optional Q/K/V prefix on its local masked FA path.
• Use dense FA3 scheduling for a single packed sequence, then scatter back into the BCG bucket.
• Route Qwen-Image's masked TP path through segmented packing.
• Add registered H100/H200 kernel correctness coverage and a production-shape benchmark.
• Document/export the new diffusion data-movement kernel.

Accuracy Tests

Validated on the latest main used for this PR (b647ae8) and 2x NVIDIA H200:

• pytest -q test/registered/kernels/ops/diffusion/test_layout.py -k "segmented_pack or dense_scheduler": 16 passed.
• Segmented packing is bitwise equal to materialized torch.cat + fused_pack_qkv across BF16/FP16 and batch/shape variants.
• Five Qwen-Image-2512 1024x1024, seed-42 model outputs were byte-identical to the baseline. SHA256 for every measured output: 862fd0c4fbaf443a5595058255e665ff0bd801f7e7cafc1b1f115a16df10161b.

Speed Tests and Profiling

Environment: 2x H200, TP=2, Qwen/Qwen-Image-2512, 1024x1024, 50 steps, guidance=1, BCG enabled, torch.compile disabled. The incremental A/B used the same CustomAllReduceV2 stack on both sides.

• Kernel microbenchmark, Qwen production shape [B=1, text=64, image=4096, H=12, D=128]: materialized 47.74 us, segmented 21.11 us (2.26x faster).
• Steady denoise: baseline 78.5-78.8 ms/step; segmented pack 77.1-77.5 ms/step; segmented pack + static FA 77.0-77.3 ms/step.
• Torch profiler: 1,212 CatArrayBatchedCopy calls / 10.768 ms disappeared from a 402-attention-call trace. The replacement segmented pack costs 18.0 us/call, matching the old pack itself.
• FA3 moved from VarlenDynamicPersistent 159.4 us plus prepare_varlen_num_blocks 3.44 us to StaticPersistent 147.1 us, matching vLLM-Omni's measured 147.4 us.

Checklist

• Format your code according to the contribution guide.
• Add unit tests according to the contribution guide.
• Update documentation according to the contribution guide.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

Background benchmark investigation: BBuf/how-to-optim-algorithm-in-cuda#21

---

CI States

Latest PR Test (Base): ❌ Run #33069117658
Latest PR Test (Extra): ❌ Run #33069117413
Latest PR Test (AMD ROCm 7.2): ❌ Run #33069117562

[BBuf]

Consolidated into #36680 as requested. The full segmented-pack/static-FA commit, benchmark, tests, and docs are now preserved there.