Repository navigation
Conversation
…ton launch Add fused_joint_qkv_norm_rope: a two-stream variant of the fused QK RMSNorm+RoPE kernel that normalizes text and image Q/K with their own weights and writes the rotated Q/K rows and the V rows straight into the joint [B, S_txt + S_img, H, D] tensors attention consumes, replacing the double block's four norms, three cat copies and two RoPE passes with one launch and no intermediate copies. Single blocks reuse the existing single-stream op over the already-joint sequence. One packed rope table is built per forward (activation dtype, batch-expanded) and passed through joint_attention_kwargs; SP keeps the eager per-stream chain. Measured on one H200 (vLLM 0.28 host tree, see PR): per-block prep -80..-89%, full-depth Flux.2 transformer forward -8.6% at 1024^2. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
…rm + RoPE path The klein transformer is a copy of the Flux.2 one with the same double/ single block structure; wire it to the same ops and per-forward packed table so the fused path (and its end-to-end measurement on the ungated FLUX.2-klein checkpoints) covers both pipelines. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
… module pack_qk_norm_rope_table() now lives next to the ops so other dual-stream consumers (FLUX.1, Ovis-Image, ...) reuse it; Flux.2/klein keep their thin wrapper. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
… packing on kernel availability, never mutate caller kwargs - _qk_norm_rope_kernel now calls _norm_rope_tile like the two-stream kernel (one copy of the RMSNorm/RoPE arithmetic; existing bitwise tests cover it). - pack_qk_norm_rope_table returns None (no allocation, no copy) unless fused_qk_norm_rope_available() says the CUDA kernel would run for the activation dtype/device/geometry — CPU, NPU, ROCm and non-bf16 paths pay nothing. - Flux.2/klein rebind joint_attention_kwargs to a copy instead of writing the per-forward table into the caller-owned dict. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
…over the gate in tests - QK_NORM_ROPE_TABLE_KEY lives next to pack_qk_norm_rope_table; flux2_klein no longer imports Flux.2's private helpers. - pack_qk_norm_rope_table takes sequence_parallel_size and returns None under SP, so the model forwards call it directly. - Tests: token-gate boundary and env override (CPU), sequence-parallel skip (CPU), and both ops bitwise-equal to eager under torch.compile(fullgraph) and CUDA-graph capture/replay. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
A trailing stream whose tokens are normalised but not rotated (HunyuanVideo's text tokens) can now share the packer instead of building its own cos/sin; the rows count toward the token gate and are exact in the kernels. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Route ZImageAttention (noise refiner, context refiner and the 30 main blocks) through the shared single-stream fused QK RMSNorm + RoPE op with one packed [cos | sin] table per attention site per forward; SP keeps the eager chain. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
… under the vLLM config Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
The refiner sites are not parallelized (_sp_plan shards only unified_prepare's outputs) and at the unified site the packer runs after that sharding, so the cos/sin it sees are already this rank's shard — the coefficients the eager chain would use on this rank. Drop the blanket SP skip and cover the case. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: hyw <yuweih205@gmail.com>
Preserve the upstream large-storage-offset fix and extend 64-bit addressing to the joint QKV kernel. Retain all existing fusion and upstream regression tests, add joint stream coverage, and satisfy the current local pre-commit checks. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Merge the updated vllm-project#7560 dependency, preserving this model integration and its numerical and sequence-parallel contracts. Include the shared 64-bit storage-offset fix and regression coverage. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
…ge/a3-on-7597-20260925
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Scope
Validate the fused QK RMSNorm + RoPE path from #7597 on Ming-Image's inherited Z-Image transformer, including direct conditioning and a reference latent. This branch is stacked on #7597 and has been synchronized with current
main; the Ming-specific change is one CUDA regression test (git diff ee9f9240522c1956ad458db259b8a87cbdb1e116 -- tests/diffusion/models/ming_image/).Part of #8134, A3. Merge after #7597.
Test result
RTX 5090, Python 3.12.13, torch 2.13.0+cu130, vLLM 0.29.0, 0z5a environment. The Ming CUDA regression test passed (1 passed) and counted fused operator calls only when the gate was enabled. Fused and eager outputs matched within BF16 tolerance.
The following is a small randomly initialized Ming transformer with one main layer and one refiner layer, 16×16 latent input, batch 1. Each row is the median of 30 CUDA-event timings after warmup; the two rows repeat the eager/fused sequence. This is an operator integration measurement, not full checkpoint latency.
Full checkpoint E2E image quality and latency are pending the model download and will be added to this draft.