Skip to content

Qwen image latent pack unpack - #410

Open
nodeeeeee wants to merge 2 commits into
RL-Align:test-qwenimagefrom
nodeeeeee:qwen-image-latent-pack-unpack
Open

nodeeeeee wants to merge 2 commits into
RL-Align:test-qwenimagefrom
nodeeeeee:qwen-image-latent-pack-unpack

Conversation

@nodeeeeee

@nodeeeeee nodeeeeee commented Sep 12, 2026 •

Copy link
Copy Markdown

Add Qwen-Image latent pack/unpack kernels

Adds CUDA, Triton, and PyTorch backends with forward/backward support for FP32, FP16, and BF16. Pack converts spatial latents into tokens by grouping 2x2 regions across channels; unpack restores the spatial layout. Both preserve values and bit patterns exactly.

CUDA uses a two-stage shared-memory transpose with a 32x32 data tile and 32x8 threads. Autograd executes the inverse kernel for backward. Backend registration preserves existing dispatch behavior.

Validation passed 241 tests. Numerical correctness was verified against the PyTorch reference for pack/unpack outputs and gradients using zero tolerances and byte-for-byte equality. Batch invariance was tested by comparing the first sample processed alone with the same sample processed in batches of 1, 2, 4, and 8. Tests also cover round trips, edge shapes, special values, CUDA streams, fallback, and build configuration. Recorded sanitizer runs reported zero memory errors and race hazards; three launch configurations passed 90 cases.

H100 measurements cover 36 configurations: C=16, three image sizes, batches 1/2/4/8, and three dtypes. The table shows geometric-mean speedups over PyTorch; values above 1 are faster. CUDA Graph ranges show minimum-to-maximum speedups. Eager backward remains slightly slower. Measurements predate the readability refactor and do not represent end-to-end training speed.

Operation CUDA Graph CUDA Graph range CUDA eager events CUDA eager wall time Triton Graph
Pack forward 1.337x 1.050-1.609x 1.092x 1.091x 0.543x
Unpack forward 1.646x 1.294-1.975x 1.096x 1.095x 0.958x
Pack backward 1.644x 1.299-1.972x 0.945x 0.945x 0.957x
Unpack backward 1.337x 1.046-1.606x 0.956x 0.956x 0.543x

#386

@coderabbitai

coderabbitai Bot commented Sep 12, 2026 •

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 977df318-047b-44d9-881a-1e9db1194dee

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
  • Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.


Comment @coderabbitai help to get the list of available commands.

@Flink-ddd Flink-ddd added multimodal Features, bugs, or optimizations specific to multimodal support. platform: cuda Specific optimizations or bugs in NVIDIA graphics cards (such as FlashInfer, TMA optimizations) labels Sep 12, 2026
@nodeeeeee
nodeeeeee changed the base branch from test to test-qwenimage October 1, 2026 11:26

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

multimodal Features, bugs, or optimizations specific to multimodal support. platform: cuda Specific optimizations or bugs in NVIDIA graphics cards (such as FlashInfer, TMA optimizations)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants