diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 4333c9934..39980c6e6 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -95,7 +95,7 @@ th {
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
+| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index c0b29cf6e..0830a280d 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -138,7 +138,7 @@ The following tables show which models support each feature:
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
diff --git a/recipes/README.md b/recipes/README.md
index f82a5955d..c5f4166ce 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -78,6 +78,7 @@ recipes/
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
+| [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
diff --git a/recipes/SenseNova/SenseNova-U1.5.md b/recipes/SenseNova/SenseNova-U1.5.md
new file mode 100644
index 000000000..80b40ab2a
--- /dev/null
+++ b/recipes/SenseNova/SenseNova-U1.5.md
@@ -0,0 +1,159 @@
+# SenseNova-U1.5
+
+> Unified image generation and understanding, with an official 8-step distilled LoRA
+
+## Summary
+
+- Vendor: SenseNova
+- Model: `sensenova/SenseNova-U1.5-8B-MoT`
+- LoRA: `sensenova/SenseNova-U1.5-8B-MoT-LoRAs` (`SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`)
+- Task: text2img, img2img, img2text (visual understanding), text2text (chat)
+- Mode: Offline inference, Online serving (OpenAI-compatible API)
+- Maintainer: Community
+
+## When to use this recipe
+
+U1.5 runs on the same `SenseNovaU1Pipeline` as U1 — the checkpoint keeps
+`model_type: neo_chat`, so it resolves without `--model-class-name`. Relative to U1 the
+config flips two fields, `use_pixel_head` to `true` (the flow-matching head becomes a
+`ConvDecoder`) and `noise_scale_max_value` from 8.0 to 16.0; both are already read by
+`SenseNovaU1Config`. The checkpoint is 13 shards / 50.2 GB on disk, 30.3 GB of that fp32, and
+loads to roughly 34 GB in bf16.
+
+Use this recipe for U1.5 specifically, including its distilled few-step LoRA. For U1 see
+[SenseNova-U1](SenseNova-U1.md).
+
+## Hardware Support
+
+### GPU
+
+#### 1x A800 80GB
+
+- 1x NVIDIA A800 80GB PCIe with an Intel Xeon Gold 6336Y, driver 595.84, Linux
+- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
+ transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
+- Peak GPU memory: 34.3 GB at 1024x1024, 36.4 GB at 1536x2720
+
+##### Text-to-Image
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+ --model sensenova/SenseNova-U1.5-8B-MoT \
+ --prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
+ --width 1024 --height 1024 \
+ --seed 42 --num-inference-steps 50 --cfg-scale 4.0 \
+ --extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
+ --output sensenova_u15_t2i.png
+```
+
+Think mode (`"think": true`) is recommended for higher image quality.
+
+##### Image-to-Image Editing
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+ --model sensenova/SenseNova-U1.5-8B-MoT \
+ --prompt "Turn this into an oil painting" \
+ --image input.png --resolution 1024 \
+ --seed 42 --num-inference-steps 25 --cfg-scale 4.0 \
+ --extra-args '{"think": true, "img_cfg_scale": 1.0, "cfg_norm": "none", "timestep_shift": 3.0}' \
+ --output sensenova_u15_edit.png
+```
+
+##### 8-step distilled LoRA
+
+The distilled LoRA is fused into the generation tower at load time:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+ --model sensenova/SenseNova-U1.5-8B-MoT \
+ --lora-path SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors --lora-backend distill \
+ --prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
+ --width 1024 --height 1024 \
+ --seed 42 --num-inference-steps 8 --cfg-scale 1.0 \
+ --extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
+ --output sensenova_u15_lora8.png
+```
+
+**Use `--cfg-scale 1.0` with this LoRA.** It is distilled with DMD and runs without
+classifier-free guidance; the default `4.0` applies guidance twice and produces a
+blown-out, posterised image.
+
+##### Online serving
+
+```bash
+vllm serve sensenova/SenseNova-U1.5-8B-MoT --omni --port 8091
+
+python examples/online_serving/sensenova_u1/openai_chat_client.py \
+ -s http://127.0.0.1:8091 -m img2text -i input.png -p "Describe this image."
+```
+
+`-s` takes the base URL; the client appends `/v1` itself.
+
+#### Measured latency (1x A800 80GB, 25 steps, median of 3 after a warmup)
+
+| Resolution | Step latency | Total |
+| --- | --- | --- |
+| 1024x1024 | 209.5 ms | 5.24 s |
+| 1536x1536 | 458.5 ms | 11.46 s |
+
+At 1536x2720 with 50 steps, end-to-end is 43.0 s.
+
+#### Measured latency (1x NVIDIA H200 139GB, single run, no warmup)
+
+Reported by @hsliuustc0106 against PR head `29c090d`, on one reserved CUDA device.
+
+- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
+ transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
+- `sensenova/SenseNova-U1.5-8B-MoT` with `SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`,
+ seed 42, 1024x1024
+
+| Case | Stage latency | Peak GPU memory |
+| --- | --- | --- |
+| 50 steps, cfg 4.0, think off | 9017 ms | 34384 MiB |
+| 8 steps, cfg 1.0, no LoRA control | 624.00 ms, 77.85 ms/step | 34376 MiB |
+| 8 steps, cfg 1.0, distilled LoRA | 623.85 ms, 77.98 ms/step | 34364 MiB |
+| 50 steps, cfg 4.0, think on, `VLLM_OMNI_SENSENOVA_PAGED_DECODE=1` | 81622 ms | 34594 MiB |
+
+Each row runs the Text-to-Image command above with `--num-inference-steps` and
+`--cfg-scale` as listed and `think` set through `--extra-body`; the LoRA row adds
+`--lora-path` and `--lora-backend distill`.
+
+Single runs without a preceding warmup, so the think-on row carries the first-request
+compile rather than a steady state. The distilled LoRA fused into 168 parameters, and its
+image differs from the same-seed no-LoRA control in 1,048,574 of 1,048,576 pixels, MAE 58.62.
+
+```bash
+pytest -m "cpu and not cuda" tests/diffusion/models/sensenova_u1/ tests/diffusion/lora/test_loader.py -q
+pytest -m "cpu and not cuda" tests/config/test_environment_variables.py tests/diffusion/test_diffusion_worker.py -q
+pytest -m cuda tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py \
+ tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py -q --tb=short
+```
+
+71, 13 and 9 passed. Serving reached `/health`; `img2text` returned a 1921-character
+description and `text2text` the expected answer.
+
+#### Verification
+
+```bash
+pytest -q tests/diffusion/models/sensenova_u1/
+```
+
+## Notes
+
+- Both `use_pixel_head` and `noise_scale_max_value` come from `config.json`; no flag is needed.
+- The checkpoint carries no `configuration_neo_chat.py`, so the loader logs a few
+ "does not appear to have a file named configuration_neo_chat.py" errors during startup and
+ then proceeds on the in-tree config. Generation is unaffected.
+- The LoRA targets the generation tower only (`*_mot_gen`); understanding-tower weights are
+ untouched.
+- Autoregressive decode (think, text-to-text and image-to-text) runs on a paged K/V cache under
+ a captured CUDA graph. It falls back to the ordinary cache when the device or the bundled
+ `flash_attn_varlen_func` cannot support it; set `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` to force
+ that fallback.
+- The first request after startup costs about 0.7 s more than the steady state whether the paged
+ path is on or off. Measured on one A800 with the inductor, triton and vLLM compile caches all
+ cleared, median of three runs: 718 ms above steady with the path on, 679 ms with it off.
+- Each request captures its own graphs, and a think request captures twice because the sequence
+ grows past the 512 bucket, so the capture cost is paid per request rather than once at
+ startup.
diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py
index 00e0dffdb..9b654a7de 100644
--- a/tests/config/test_environment_variables.py
+++ b/tests/config/test_environment_variables.py
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
- EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
+ EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2,
}
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
)
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
- ModelEnvironmentVariableDisposition.PROMOTE: 33,
+ ModelEnvironmentVariableDisposition.PROMOTE: 34,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py
index d2214f940..f13c0f6c3 100644
--- a/tests/diffusion/lora/test_loader.py
+++ b/tests/diffusion/lora/test_loader.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from collections import defaultdict
import pytest
@@ -106,7 +109,7 @@ def make_dummy_lora_state_dict(
"""
key_a = f"{base_key}.{lora_a_suffix}"
key_b = f"{base_key}.{lora_b_suffix}"
- key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
+ key_bias = f"{base_key}.{lora_bias_suffix}"
if weight_a is None:
weight_a = torch.randn(RANK, HEAD_DIM)
@@ -436,3 +439,54 @@ def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
assert_params_equal(pipeline.transformer, original_high)
assert_params_equal(pipeline.transformer_2, original_low)
assert "adapter0" not in pipeline.lora_loaded
+
+
+def test_prepare_lora_delta_disambiguates_prefix_overlapping_stacked_names():
+ """A fused name can be a prefix of another fused name.
+
+ SenseNova-U1 carries both `.qkv_proj` (understanding) and `.qkv_proj_mot_gen`
+ (generation). Selecting stacked mappings by substring matched both rules for
+ the generation weight and concatenated six shards into a delta twice as tall
+ as the parameter.
+ """
+ param_to_weight_names = {
+ ".qkv_proj_mot_gen": [".q_proj_mot_gen", ".k_proj_mot_gen", ".v_proj_mot_gen"],
+ ".qkv_proj": [".q_proj", ".k_proj", ".v_proj"],
+ }
+ state_dict = {}
+ for part in ("q", "k", "v"):
+ for tower in ("_mot_gen", ""):
+ base = f"layers.0.self_attn.{part}_proj{tower}"
+ state_dict[f"{base}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
+ state_dict[f"{base}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
+
+ for base_key in ("layers.0.self_attn.qkv_proj_mot_gen", "layers.0.self_attn.qkv_proj"):
+ delta, _ = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
+ assert delta.shape == (3 * HEAD_DIM, HEAD_DIM), base_key
+
+
+@pytest.mark.parametrize(
+ "fused, shards",
+ [
+ # qwen_image_transformer.py stacked_params_mapping
+ (".to_qkv", (".to_q", ".to_k", ".to_v")),
+ (".add_kv_proj", (".add_q_proj", ".add_k_proj", ".add_v_proj")),
+ # wan2_2_transformer.py stacked_params_mapping
+ (".attn1.to_qkv", (".attn1.to_q", ".attn1.to_k", ".attn1.to_v")),
+ # wan2_2_s2v_transformer.py stacked_params_mapping
+ (".self_attn.to_qkv", (".self_attn.q", ".self_attn.k", ".self_attn.v")),
+ ],
+)
+def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards):
+ """None of the fused names shipped today overlap by prefix, so matching on the
+ tail must select exactly the shards substring matching selected."""
+ param_to_weight_names = {fused: list(shards)}
+ base_key = f"blocks.0{fused}"
+ state_dict = {}
+ for shard in shards:
+ state_dict[f"blocks.0{shard}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
+ state_dict[f"blocks.0{shard}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
+
+ delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
+ assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM)
+ assert len(used_keys) == 2 * len(shards)
diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py
new file mode 100644
index 000000000..e628a5224
--- /dev/null
+++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py
@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for SenseNova-U1 attention head handling.
+
+The model used to expand K/V from 8 heads to 32 with ``repeat_interleave``
+before calling the attention backend, so the backend's own GQA dispatch never
+saw a GQA shape, and it passed an all-zeros mask during single-token decode,
+which only served to keep SDPA off its fused kernels. These pin both contracts:
+K/V must reach the backend uncompressed, decode must carry no mask, and the
+compressed path must compute the same thing the expanded one did.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+ SenseNovaU1Attention,
+ SenseNovaU1Model,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+N_HEADS, N_KV_HEADS, HEAD_DIM, SEQ = 32, 8, 8, 6
+GROUPS = N_HEADS // N_KV_HEADS
+
+
+class _Recorder:
+ """Stands in for the Attention module and records what it was handed."""
+
+ def __init__(self):
+ self.calls: list[tuple[torch.Size, torch.Size, torch.Tensor | None]] = []
+
+ def __call__(self, query, key, value, attn_metadata):
+ mask = None if attn_metadata is None else attn_metadata.attn_mask
+ self.calls.append((query.shape, key.shape, mask))
+ return query
+
+
+class _AttnHost:
+ """Carries only the attributes ``_run_attn``/``_run_attn_bshd`` touch."""
+
+ num_kv_groups = GROUPS
+ _align_mask_dtype = staticmethod(SenseNovaU1Attention._align_mask_dtype)
+ _run_attn = SenseNovaU1Attention._run_attn
+ _run_attn_bshd = SenseNovaU1Attention._run_attn_bshd
+
+ def __init__(self):
+ self.attn = _Recorder()
+
+
+def test_bhsd_path_hands_the_backend_compressed_kv():
+ host = _AttnHost()
+ q = torch.randn(1, N_HEADS, SEQ, HEAD_DIM)
+ k = torch.randn(1, N_KV_HEADS, SEQ, HEAD_DIM)
+ host._run_attn(q, k, k.clone(), None)
+ _, key_shape, _ = host.attn.calls[-1]
+ # [B, S, H, D] after the transpose, so heads sit on dim 2.
+ assert key_shape[2] == N_KV_HEADS, f"K reached the backend with {key_shape[2]} heads"
+
+
+def test_bshd_path_hands_the_backend_compressed_kv():
+ host = _AttnHost()
+ q = torch.randn(1, SEQ, N_HEADS, HEAD_DIM)
+ k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
+ host._run_attn_bshd(q, k, k.clone(), None)
+ _, key_shape, _ = host.attn.calls[-1]
+ assert key_shape[2] == N_KV_HEADS, f"K reached the backend with {key_shape[2]} heads"
+
+
+def test_compressed_kv_computes_what_expansion_computed():
+ """The backend must return the same tensor either way, or this is a bugfix
+ dressed up as a perf change."""
+ torch.manual_seed(0)
+ q = torch.randn(1, SEQ, N_HEADS, HEAD_DIM)
+ k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
+ v = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM)
+ impl = SDPAImpl(
+ num_heads=N_HEADS,
+ head_size=HEAD_DIM,
+ softmax_scale=HEAD_DIM**-0.5,
+ num_kv_heads=N_KV_HEADS,
+ )
+ compressed = impl._forward_impl(q, k, v)
+ expanded = impl._forward_impl(q, k.repeat_interleave(GROUPS, dim=2), v.repeat_interleave(GROUPS, dim=2))
+ torch.testing.assert_close(compressed, expanded, atol=1e-6, rtol=1e-6)
+
+
+class _MaskProbe(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.masks: list[torch.Tensor | None] = []
+
+ def forward(self, hidden_states, **kwargs):
+ self.masks.append(kwargs["attention_mask"]["full_attention"])
+ return hidden_states
+
+
+def _fake_rope(hidden_states, position_ids):
+ """The forward builds the three RoPE tables before the layer loop; the probe
+ ignores them, so any 2-tuple will do."""
+ return hidden_states, hidden_states
+
+
+class _ModelHost:
+ """Carries only the attributes ``SenseNovaU1Model.forward`` touches."""
+
+ forward = SenseNovaU1Model.forward
+
+ def __init__(self, probe):
+ self.layers = [probe]
+ self.norm = torch.nn.Identity()
+ self.norm_mot_gen = torch.nn.Identity()
+ self.rotary_emb = _fake_rope
+ self.rotary_emb_hw = _fake_rope
+
+
+def _run_model(seq_len: int):
+ probe = _MaskProbe()
+ _ModelHost(probe).forward(
+ inputs_embeds=torch.zeros(1, seq_len, 4),
+ indexes=torch.zeros(3, seq_len, dtype=torch.long),
+ )
+ return probe.masks[-1]
+
+
+def test_single_token_decode_carries_no_mask():
+ assert _run_model(1) is None, "decode still builds the all-zeros mask"
+
+
+def test_prefill_still_gets_a_causal_mask():
+ mask = _run_model(3)
+ assert mask is not None and mask.shape == (1, 1, 3, 3)
+ assert mask[0, 0, 0, 0] == 0.0
+ assert torch.isinf(mask[0, 0, 0, 1]) and mask[0, 0, 0, 1] < 0
+ assert mask[0, 0, 2, 1] == 0.0
+
+
+# ---------------------------------------------------------------------------
+# The decode path changes which SDPA kernel runs, so it is not bit-identical to
+# what it replaced. "Different" would be a regression if the new result were
+# worse, so pin the direction against an independent float64 reference. This
+# needs a fused GQA kernel, which only exists on CUDA.
+# ---------------------------------------------------------------------------
+
+
+def _float64_attention(q, k, v):
+ q64 = q.double()
+ k64 = k.double().repeat_interleave(GROUPS, dim=1)
+ v64 = v.double().repeat_interleave(GROUPS, dim=1)
+ scores = (q64 @ k64.transpose(-1, -2)) * (q.shape[-1] ** -0.5)
+ return torch.softmax(scores, dim=-1) @ v64
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fused GQA")
+@pytest.mark.cuda
+@pytest.mark.L4
+@pytest.mark.parametrize("kv_len", [271, 512, 2048])
+def test_decode_without_the_no_op_mask_is_closer_to_float64(kv_len):
+ """The old decode path expanded K/V and passed an all-zeros mask, which kept
+ SDPA off its fused kernels. Dropping the mask is only defensible if the
+ kernel it unlocks is at least as accurate."""
+ dev, dtype = "cuda", torch.bfloat16
+ heads, dim = 32, 128
+ new_errs: list[float] = []
+ old_errs: list[float] = []
+ trials = 8
+ for t in range(trials):
+ torch.manual_seed(2000 + t)
+ q = torch.randn(1, heads, 1, dim, device=dev, dtype=dtype)
+ k = torch.randn(1, heads // GROUPS, kv_len, dim, device=dev, dtype=dtype)
+ v = torch.randn(1, heads // GROUPS, kv_len, dim, device=dev, dtype=dtype)
+ ref = _float64_attention(q, k, v)
+
+ zeros_mask = torch.zeros(1, 1, 1, kv_len, device=dev, dtype=dtype)
+ old = torch.nn.functional.scaled_dot_product_attention(
+ q,
+ k.repeat_interleave(GROUPS, dim=1),
+ v.repeat_interleave(GROUPS, dim=1),
+ attn_mask=zeros_mask,
+ )
+ new = torch.nn.functional.scaled_dot_product_attention(q, k, v, enable_gqa=True)
+
+ new_errs.append((new.double() - ref).abs().mean().item())
+ old_errs.append((old.double() - ref).abs().mean().item())
+
+ # A single trial can tie either way depending on the GPU and the CUDA build,
+ # so gate on the aggregate: the maskless path must not be worse on average.
+ new_mean = sum(new_errs) / trials
+ old_mean = sum(old_errs) / trials
+ wins = sum(n <= o for n, o in zip(new_errs, old_errs))
+ assert new_mean <= old_mean, (
+ f"kv_len={kv_len}: maskless mean error {new_mean:.4e} is worse than masked {old_mean:.4e} "
+ f"({wins}/{trials} per-trial wins)"
+ )
diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_distill_lora.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_distill_lora.py
new file mode 100644
index 000000000..cd66e7a0d
--- /dev/null
+++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_distill_lora.py
@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for SenseNova-U1 distilled-LoRA fusion.
+
+The checkpoints ship unfused ``q_proj_mot_gen`` / ``gate_proj`` keys while the
+pipeline runs fused layers, so each delta lands in its own row slice.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+from safetensors.torch import save_file
+
+from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+HIDDEN, KV, INTER, RANK, ALPHA = 8, 4, 16, 2, 4.0
+
+
+class _Layer(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.self_attn = nn.Module()
+ self.self_attn.qkv_proj_mot_gen = nn.Linear(HIDDEN, HIDDEN + 2 * KV, bias=False)
+ self.self_attn.o_proj_mot_gen = nn.Linear(HIDDEN, HIDDEN, bias=False)
+ self.mlp_mot_gen = nn.Module()
+ self.mlp_mot_gen.gate_up_proj = nn.Linear(HIDDEN, 2 * INTER, bias=False)
+ self.mlp_mot_gen.down_proj = nn.Linear(INTER, HIDDEN, bias=False)
+
+
+class _Stub(SenseNovaU1Pipeline):
+ """Only the parameter tree matters here, so skip the real __init__."""
+
+ def __init__(self):
+ nn.Module.__init__(self)
+ self.language_model = nn.Module()
+ self.language_model.model = nn.Module()
+ self.language_model.model.layers = nn.ModuleList([_Layer()])
+
+
+def _triple(out: int, in_: int, fill: float) -> dict[str, torch.Tensor]:
+ return {
+ "lora_down.weight": torch.full((RANK, in_), fill, dtype=torch.float32),
+ "lora_up.weight": torch.full((out, RANK), 1.0, dtype=torch.float32),
+ "alpha": torch.tensor(ALPHA),
+ }
+
+
+def _write(tmp_path, entries: dict[str, dict[str, torch.Tensor]]):
+ flat = {f"{mod}.{leaf}": t for mod, parts in entries.items() for leaf, t in parts.items()}
+ path = tmp_path / "lora.safetensors"
+ save_file(flat, str(path))
+ return str(path)
+
+
+def _prefix(i: int = 0) -> str:
+ return f"language_model.model.layers.{i}"
+
+
+def test_qkv_shards_land_in_their_own_row_slices(tmp_path):
+ pipe = _Stub()
+ qkv = pipe.language_model.model.layers[0].self_attn.qkv_proj_mot_gen.weight
+ with torch.no_grad():
+ qkv.zero_()
+ entries = {
+ f"{_prefix()}.self_attn.q_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0),
+ f"{_prefix()}.self_attn.k_proj_mot_gen": _triple(KV, HIDDEN, 2.0),
+ f"{_prefix()}.self_attn.v_proj_mot_gen": _triple(KV, HIDDEN, 3.0),
+ }
+ pipe.load_lora_weights(_write(tmp_path, entries))
+
+ scale = ALPHA / RANK
+ assert torch.allclose(qkv[:HIDDEN], torch.full((HIDDEN, HIDDEN), scale * RANK * 1.0))
+ assert torch.allclose(qkv[HIDDEN : HIDDEN + KV], torch.full((KV, HIDDEN), scale * RANK * 2.0))
+ assert torch.allclose(qkv[HIDDEN + KV :], torch.full((KV, HIDDEN), scale * RANK * 3.0))
+
+
+def test_gate_and_up_land_in_their_own_halves(tmp_path):
+ pipe = _Stub()
+ gu = pipe.language_model.model.layers[0].mlp_mot_gen.gate_up_proj.weight
+ with torch.no_grad():
+ gu.zero_()
+ entries = {
+ f"{_prefix()}.mlp_mot_gen.gate_proj": _triple(INTER, HIDDEN, 1.0),
+ f"{_prefix()}.mlp_mot_gen.up_proj": _triple(INTER, HIDDEN, 5.0),
+ }
+ pipe.load_lora_weights(_write(tmp_path, entries))
+
+ scale = ALPHA / RANK
+ assert torch.allclose(gu[:INTER], torch.full((INTER, HIDDEN), scale * RANK * 1.0))
+ assert torch.allclose(gu[INTER:], torch.full((INTER, HIDDEN), scale * RANK * 5.0))
+
+
+def test_unfused_targets_are_added_whole(tmp_path):
+ pipe = _Stub()
+ o = pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight
+ before = o.detach().clone()
+ entries = {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)}
+ pipe.load_lora_weights(_write(tmp_path, entries))
+ assert torch.allclose(o - before, torch.full((HIDDEN, HIDDEN), (ALPHA / RANK) * RANK))
+
+
+def test_lora_that_matches_nothing_raises_instead_of_no_op(tmp_path):
+ """Without a loader the worker only warned and generated without the
+ adapter, so a wrong path produced a plausible image and no error."""
+ pipe = _Stub()
+ entries = {f"{_prefix()}.self_attn.does_not_exist": _triple(HIDDEN, HIDDEN, 1.0)}
+ with pytest.raises(ValueError, match="matched no parameter"):
+ pipe.load_lora_weights(_write(tmp_path, entries))
+
+
+def test_shape_mismatch_raises(tmp_path):
+ pipe = _Stub()
+ entries = {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN + 1, HIDDEN, 1.0)}
+ with pytest.raises(RuntimeError, match="must match the size"):
+ pipe.load_lora_weights(_write(tmp_path, entries))
+
+
+def test_multiple_files_rejected(tmp_path):
+ pipe = _Stub()
+ path = _write(tmp_path, {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)})
+ with pytest.raises(ValueError, match="exactly one"):
+ pipe.load_lora_weights([path, path])
+
+
+class _ShardedLayer(nn.Module):
+ """One rank of a two-way sharded fused QKV, with a vLLM-style weight loader.
+
+ The parameter holds half of each of q, k and v. The loader is handed the whole
+ layer's tensor and keeps only this rank's slice, which is what vLLM does.
+ """
+
+ def __init__(self):
+ super().__init__()
+ self.weight = nn.Parameter(torch.zeros(HIDDEN // 2 + KV, HIDDEN))
+ self.seen: list[tuple[str, tuple[int, ...]]] = []
+ self.weight.weight_loader = self._weight_loader
+
+ def _weight_loader(self, param, loaded_weight, shard_id):
+ self.seen.append((shard_id, tuple(loaded_weight.shape)))
+ sizes = {"q": HIDDEN // 2, "k": KV // 2, "v": KV // 2}
+ offsets = {"q": 0, "k": HIDDEN // 2, "v": HIDDEN // 2 + KV // 2}
+ rows = sizes[shard_id]
+ param.data[offsets[shard_id] : offsets[shard_id] + rows] = loaded_weight[:rows]
+
+
+class _ShardedStub(SenseNovaU1Pipeline):
+ def __init__(self):
+ nn.Module.__init__(self)
+ self.language_model = nn.Module()
+ self.language_model.model = nn.Module()
+ layer = nn.Module()
+ layer.self_attn = nn.Module()
+ layer.self_attn.qkv_proj_mot_gen = _ShardedLayer()
+ self.language_model.model.layers = nn.ModuleList([layer])
+
+
+def test_fused_delta_is_sharded_through_the_weight_loader(tmp_path):
+ """Under tensor parallelism the parameter is a rank-local shard while the
+ checkpoint delta spans the whole layer, so the delta has to go through the
+ layer's weight loader. Adding it directly raises a shape error."""
+ pipe = _ShardedStub()
+ layer = pipe.language_model.model.layers[0].self_attn.qkv_proj_mot_gen
+ entries = {
+ f"{_prefix()}.self_attn.q_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0),
+ f"{_prefix()}.self_attn.k_proj_mot_gen": _triple(KV, HIDDEN, 2.0),
+ f"{_prefix()}.self_attn.v_proj_mot_gen": _triple(KV, HIDDEN, 3.0),
+ }
+ pipe.load_lora_weights(_write(tmp_path, entries))
+
+ # the loader saw the whole layer's tensors, not a pre-sharded slice
+ assert layer.seen == [("q", (HIDDEN, HIDDEN)), ("k", (KV, HIDDEN)), ("v", (KV, HIDDEN))]
+
+ scale = ALPHA / RANK * RANK
+ got = layer.weight
+ assert torch.allclose(got[: HIDDEN // 2], torch.full((HIDDEN // 2, HIDDEN), scale * 1.0))
+ assert torch.allclose(got[HIDDEN // 2 : HIDDEN // 2 + KV // 2], torch.full((KV // 2, HIDDEN), scale * 2.0))
+ assert torch.allclose(got[HIDDEN // 2 + KV // 2 :], torch.full((KV // 2, HIDDEN), scale * 3.0))
+
+
+def test_fusion_does_not_retain_the_adapter_tensors(tmp_path):
+ """Fusion is one-way and this pipeline has no unload path, so holding the
+ fp32 state dict would keep about 1.5 GiB alive for the shipped checkpoint."""
+ pipe = _Stub()
+ path = _write(tmp_path, {f"{_prefix()}.self_attn.o_proj_mot_gen": _triple(HIDDEN, HIDDEN, 1.0)})
+ pipe.load_lora_weights(path, adapter_name="distill")
+
+ kept = pipe.lora_loaded["distill"]
+ assert not isinstance(kept, dict)
+ assert not isinstance(kept, torch.Tensor)
+
+ # the sentinel still makes a second load a no-op
+ before = pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight.detach().clone()
+ pipe.load_lora_weights(path, adapter_name="distill")
+ torch.testing.assert_close(pipe.language_model.model.layers[0].self_attn.o_proj_mot_gen.weight, before)
diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py
new file mode 100644
index 000000000..cae40b220
--- /dev/null
+++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py
@@ -0,0 +1,800 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for the paged decode cache and its CUDA graph.
+
+The cache exists so a decode step has static shapes and can be captured. That
+puts the burden on two things a normal test would not look at: the prefix must
+survive the hand-off from the ordinary cache, and every quantity that varies
+between steps must live in a device tensor. A Python-level index is correct
+eagerly and silently wrong once captured, because capture bakes it in -- which
+is exactly what the last test here pins.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.lora.layers import BaseLayerWithLoRA
+
+import vllm_omni.diffusion.models.sensenova_u1.paged_decode as paged_decode
+from vllm_omni.diffusion.models.sensenova_u1.paged_decode import (
+ BLOCK_SIZE,
+ BUCKETS,
+ TAIL_STEP,
+ PagedDecodeCache,
+ _bucket_for,
+)
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+ SenseNovaU1Attention,
+)
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+LAYERS, KV_HEADS, HEAD_DIM = 2, 4, 8
+N_HEADS = KV_HEADS * 2
+
+
+class _Layer:
+ def __init__(self, keys=None, values=None):
+ self.keys, self.values = keys, values
+
+
+class _Cache:
+ """Stands in for the transformers cache: layers with [B, H, S, D] tensors."""
+
+ def __init__(self, layers):
+ self.layers = layers
+
+
+def _dyn_cache(prefix, device=None):
+ torch.manual_seed(0)
+ return _Cache(
+ [
+ _Layer(
+ torch.randn(1, KV_HEADS, prefix, HEAD_DIM, device=device),
+ torch.randn(1, KV_HEADS, prefix, HEAD_DIM, device=device),
+ )
+ for _ in range(LAYERS)
+ ]
+ )
+
+
+def test_bucket_rounds_up_and_is_block_aligned():
+ assert _bucket_for(1) == 512
+ assert _bucket_for(512) == 512
+ assert _bucket_for(513) == 1024
+ assert _bucket_for(100_000) % BLOCK_SIZE == 0
+
+
+def test_past_the_last_bucket_the_schedule_grows_in_steps():
+ """Every bucket change reallocates the cache and re-captures the graph, so
+ the tail of the schedule must not step by BLOCK_SIZE. A think edit with two
+ 2048x2048 inputs runs to about 9,100 tokens, so this regime is reachable,
+ and a request decodes at most 1024 steps.
+ """
+ last = BUCKETS[-1]
+ assert _bucket_for(last + 1) == last + TAIL_STEP
+ assert _bucket_for(100_000) % BLOCK_SIZE == 0
+ for start in (last + 1, last + TAIL_STEP - 1, last + TAIL_STEP + 7):
+ bucket, grows = _bucket_for(start), 0
+ for n in range(start, start + 1025):
+ if _bucket_for(n) != bucket:
+ bucket = _bucket_for(n)
+ grows += 1
+ assert grows <= 1, f"start {start}: {grows} reallocations over 1024 decode steps"
+
+
+def test_prefix_survives_the_hand_off():
+ dyn = _dyn_cache(37)
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+ assert paged.length == 37
+ for i in range(LAYERS):
+ stored = paged.k[i].view(-1, KV_HEADS, HEAD_DIM)[:37]
+ torch.testing.assert_close(stored, dyn.layers[i].keys[0].transpose(0, 1))
+
+
+def test_write_back_restores_what_was_handed_over():
+ """Decode is paged but the stage after it reads the ordinary cache."""
+ dyn = _dyn_cache(21)
+ original = [layer.keys.clone() for layer in dyn.layers]
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+ paged.to_dynamic_cache(dyn)
+ for before, layer in zip(original, dyn.layers):
+ torch.testing.assert_close(layer.keys, before)
+
+
+def test_growing_keeps_the_tokens_already_stored():
+ dyn = _dyn_cache(500)
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+ assert paged.bucket == 512
+ before = [k.view(-1, KV_HEADS, HEAD_DIM)[:500].clone() for k in paged.k]
+ gen = paged.generation
+ assert paged.grow(513) is True
+ assert paged.bucket == 1024
+ assert paged.generation > gen, "a reallocation must invalidate captured graphs"
+ for old, new in zip(before, paged.k):
+ torch.testing.assert_close(new.view(-1, KV_HEADS, HEAD_DIM)[:500], old)
+
+
+def test_a_failed_grow_leaves_the_cache_on_its_old_buffers(monkeypatch):
+ """A half-applied resize would strand a capture on freed storage.
+
+ ``generation`` is half of the key a captured graph is stored under, so
+ publishing new K without publishing V, the block table, the bucket and the
+ generation leaves the old key selecting a graph whose recorded addresses
+ have just been freed. A later request replays it. Injecting a failure into
+ the last allocation is the cheapest way to reach that window.
+ """
+ dyn = _dyn_cache(500)
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+ before_k, before_v = list(paged.k), list(paged.v)
+ before_table, before_bucket = paged.block_table, paged.bucket
+ before_generation, before_length = paged.generation, paged.length
+
+ real_zeros = torch.zeros
+ calls = []
+
+ def flaky(*args, **kwargs):
+ calls.append(1)
+ if len(calls) == 2 * LAYERS: # the last of the two allocations per layer
+ raise RuntimeError("injected allocation failure")
+ return real_zeros(*args, **kwargs)
+
+ monkeypatch.setattr(torch, "zeros", flaky)
+ with pytest.raises(RuntimeError, match="injected"):
+ paged.grow(513)
+ monkeypatch.undo()
+
+ assert [id(t) for t in paged.k] == [id(t) for t in before_k], "K was published on its own"
+ assert [id(t) for t in paged.v] == [id(t) for t in before_v], "V was published on its own"
+ assert paged.block_table is before_table
+ assert paged.bucket == before_bucket
+ assert paged.generation == before_generation, "the graph key moved without new buffers"
+ assert paged.length == before_length
+
+
+def test_length_lives_in_device_tensors():
+ """Both the attended length and the write slot must be tensors.
+
+ A captured graph reads them at replay; a Python int would be frozen at
+ capture time and every replay would attend the wrong span and overwrite the
+ same slot.
+ """
+ dyn = _dyn_cache(10)
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, LAYERS, torch.device("cpu"), torch.float32)
+ paged.set_length(15)
+ assert torch.is_tensor(paged.seqused) and int(paged.seqused) == 15
+ assert torch.is_tensor(paged.pos) and int(paged.pos) == 14
+
+
+# ---------------------------------------------------------------------------
+# The capture contract. These need CUDA and the bundled flash-attention kernel.
+# ---------------------------------------------------------------------------
+
+cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_paged_attention_matches_sdpa_over_the_same_buffer():
+ import torch.nn.functional as F
+
+ from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+ heads, kv_heads, dim, n = 8, 2, 64, 100
+ dev, dt = torch.device("cuda"), torch.bfloat16
+ torch.manual_seed(0)
+ dyn = _Cache(
+ [
+ _Layer(
+ torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+ torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+ )
+ ]
+ )
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
+ paged.set_length(n)
+ q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+ k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
+ scale = dim**-0.5
+ # write the last slot back as itself so attend() is a pure read
+ k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
+ v = torch.zeros_like(k)
+ v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
+ got = paged.attend(0, q, k, v, scale)
+ flat_k = paged.k[0].view(-1, kv_heads, dim)[:n]
+ flat_v = paged.v[0].view(-1, kv_heads, dim)[:n]
+ want = F.scaled_dot_product_attention(
+ q,
+ flat_k.unsqueeze(0).transpose(1, 2).contiguous(),
+ flat_v.unsqueeze(0).transpose(1, 2).contiguous(),
+ enable_gqa=True,
+ scale=scale,
+ ).transpose(1, 2)
+ torch.testing.assert_close(got.float(), want.float(), atol=2e-3, rtol=2e-3)
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_captured_graph_follows_a_later_length():
+ """One capture has to serve every length in the bucket.
+
+ If the attended span were a Python value the graph would keep attending the
+ capture-time length, and the run would still look fast.
+ """
+ import torch.nn.functional as F
+
+ from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+ heads, kv_heads, dim, n = 8, 2, 64, 200
+ dev, dt = torch.device("cuda"), torch.bfloat16
+ torch.manual_seed(0)
+ dyn = _Cache(
+ [
+ _Layer(
+ torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+ torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+ )
+ ]
+ )
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
+ q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+ k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
+ v = torch.zeros_like(k)
+ scale = dim**-0.5
+
+ paged.set_length(n)
+ k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
+ v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
+
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(3):
+ paged.attend(0, q, k, v, scale)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ out = paged.attend(0, q, k, v, scale)
+
+ shorter = n - 43
+ paged.set_length(shorter)
+ graph.replay()
+ torch.accelerator.synchronize()
+
+ flat_k = paged.k[0].view(-1, kv_heads, dim)[:shorter]
+ flat_v = paged.v[0].view(-1, kv_heads, dim)[:shorter]
+ want = F.scaled_dot_product_attention(
+ q,
+ flat_k.unsqueeze(0).transpose(1, 2).contiguous(),
+ flat_v.unsqueeze(0).transpose(1, 2).contiguous(),
+ enable_gqa=True,
+ scale=scale,
+ ).transpose(1, 2)
+ torch.testing.assert_close(out.float(), want.float(), atol=2e-3, rtol=2e-3)
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_captured_graph_writes_each_step_to_its_own_slot():
+ """The write slot has to be a device tensor, not a Python index.
+
+ ``test_a_captured_graph_follows_a_later_length`` pins the attended span.
+ The slot written is the other half: a Python index is frozen at capture, so
+ every replay would overwrite the capture-time slot and the cache would stop
+ growing -- correct eagerly, silently wrong once captured.
+ """
+ from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+ heads, kv_heads, dim, n = 8, 2, 64, 100
+ dev, dt = torch.device("cuda"), torch.bfloat16
+ torch.manual_seed(0)
+ dyn = _Cache(
+ [
+ _Layer(
+ torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+ torch.randn(1, kv_heads, n, dim, device=dev, dtype=dt),
+ )
+ ]
+ )
+ paged = PagedDecodeCache.from_dynamic_cache(dyn, 1, dev, dt)
+ q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+ k = torch.zeros(1, kv_heads, 1, dim, device=dev, dtype=dt)
+ v = torch.zeros_like(k)
+ scale = dim**-0.5
+
+ # Capture on the step that writes slot n-1, restoring what is already there
+ # so the capture itself changes nothing.
+ paged.set_length(n)
+ k[0, :, 0] = dyn.layers[0].keys[0, :, n - 1]
+ v[0, :, 0] = dyn.layers[0].values[0, :, n - 1]
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(3):
+ paged.attend(0, q, k, v, scale)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ paged.attend(0, q, k, v, scale)
+
+ kept = paged.k[0].view(-1, kv_heads, dim)[n - 1].clone()
+
+ # The next decode step: one token further, different K/V.
+ paged.set_length(n + 1)
+ k[0, :, 0] = 1.5
+ v[0, :, 0] = -2.5
+ graph.replay()
+ current_omni_platform.synchronize()
+
+ flat_k = paged.k[0].view(-1, kv_heads, dim)
+ flat_v = paged.v[0].view(-1, kv_heads, dim)
+ torch.testing.assert_close(flat_k[n].float(), torch.full_like(flat_k[n], 1.5).float())
+ torch.testing.assert_close(flat_v[n].float(), torch.full_like(flat_v[n], -2.5).float())
+ torch.testing.assert_close(flat_k[n - 1].float(), kept.float())
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_whole_decode_loop_matches_the_unpaged_path():
+ """Per-step kernel equivalence is not loop equivalence.
+
+ With the switch off a single-token step attends over a cache grown by
+ ``torch.cat`` and no mask; with it on the same step reads a bucket-sized
+ buffer through ``seqused_k``. Run the same sequence of steps through both,
+ crossing a bucket boundary on the way, and require the same argmax at every
+ step off a fixed projection -- which is the token id the pipeline would emit.
+ """
+ import torch.nn.functional as F
+
+ from vllm_omni.diffusion.models.sensenova_u1.paged_decode import PagedDecodeCache
+
+ heads, kv_heads, dim = 8, 2, 64
+ prefill, steps = 500, 40 # crosses the 512 bucket at step 13
+ dev, dt = torch.device("cuda"), torch.bfloat16
+ scale = dim**-0.5
+ torch.manual_seed(7)
+
+ keys = torch.randn(1, kv_heads, prefill, dim, device=dev, dtype=dt)
+ values = torch.randn(1, kv_heads, prefill, dim, device=dev, dtype=dt)
+ paged = PagedDecodeCache.from_dynamic_cache(_Cache([_Layer(keys, values)]), 1, dev, dt)
+ ref_k, ref_v = keys.clone(), values.clone()
+ proj = torch.randn(dim * heads, 128, device=dev, dtype=torch.float32)
+
+ paged_ids, ref_ids, worst = [], [], 0.0
+ for step in range(steps):
+ q = torch.randn(1, heads, 1, dim, device=dev, dtype=dt)
+ k = torch.randn(1, kv_heads, 1, dim, device=dev, dtype=dt)
+ v = torch.randn(1, kv_heads, 1, dim, device=dev, dtype=dt)
+
+ if paged.length + 1 > paged.bucket:
+ paged.grow(paged.length + 1)
+ paged.set_length(paged.length + 1)
+ got = paged.attend(0, q, k, v, scale)
+
+ ref_k = torch.cat([ref_k, k], dim=2)
+ ref_v = torch.cat([ref_v, v], dim=2)
+ want = F.scaled_dot_product_attention(q, ref_k, ref_v, enable_gqa=True, scale=scale).transpose(1, 2)
+
+ worst = max(worst, (got.float() - want.float()).abs().max().item())
+ paged_ids.append(int((got.float().reshape(1, -1) @ proj).argmax()))
+ ref_ids.append(int((want.float().reshape(1, -1) @ proj).argmax()))
+
+ assert paged.bucket == 1024, "the loop never crossed a bucket boundary"
+ assert paged_ids == ref_ids, (
+ f"argmax diverges at step {next(i for i, (a, b) in enumerate(zip(paged_ids, ref_ids)) if a != b)}"
+ )
+ assert worst < 5e-2, f"largest per-step deviation {worst:.3e}"
+
+
+class _StubLM(torch.nn.Module):
+ """Just enough CUDA work inside the capture for a pool to be allocated."""
+
+ def __init__(self, device):
+ super().__init__()
+ self.w = torch.randn(16, 64, device=device)
+
+ def forward(self, input_ids, indexes, past_key_values=None, use_cache=False, paged_cache=None):
+ x = input_ids.float().expand(1, 16) @ self.w
+ return SimpleNamespace(logits=x.unsqueeze(0))
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_the_decode_graph_captures_into_the_shared_platform_pool(monkeypatch):
+ """``_decode_context`` builds a runner per request, so a private capture pool
+ hands every request its own arena and never gives it back. Measured on a
+ real model that was about 40 MiB of device memory per request, rising
+ linearly over twelve sequential think requests while the same twelve with
+ the switch off held flat.
+ """
+ from vllm.platforms import current_platform
+
+ from vllm_omni.diffusion.models.sensenova_u1.paged_decode import DecodeGraphRunner
+
+ pools = []
+ real_graph = torch.cuda.graph
+
+ def spy(graph, pool=None, **kwargs):
+ pools.append(pool)
+ return real_graph(graph, pool=pool, **kwargs)
+
+ monkeypatch.setattr(torch.cuda, "graph", spy)
+
+ dev = torch.device("cuda")
+ cache = PagedDecodeCache.from_dynamic_cache(_dyn_cache(8, dev), LAYERS, dev, torch.float32)
+ runner = DecodeGraphRunner(_StubLM(dev), cache, dev)
+ runner.step(0, 0)
+
+ assert pools, "no graph was captured"
+ assert pools[0] is not None, "captured into a fresh private pool"
+ assert pools[0] == current_platform.get_global_graph_pool()
+ assert not hasattr(runner, "_pool"), "the runner still keeps a pool of its own"
+
+
+# ---------------------------------------------------------------------------
+# The cache is only worth anything if production actually reaches it, and the
+# kwargs it needs are not in every wheel that exports the kernel. The tests
+# above would all still pass with the `forward_und` branch deleted, so these
+# drive the real dispatch instead of the cache object.
+# ---------------------------------------------------------------------------
+
+
+class _PagedProbe:
+ def __init__(self, out):
+ self.out = out
+ self.calls: list[tuple[int, torch.Size, float]] = []
+
+ def attend(self, layer_idx, query, key, value, scaling):
+ self.calls.append((layer_idx, query.shape, scaling))
+ return self.out
+
+
+class _AttnHost:
+ """Carries only what ``forward_und`` touches on the way to the branch."""
+
+ forward_und = SenseNovaU1Attention.forward_und
+ qkv_proj = q_norm = k_norm = q_norm_hw = k_norm_hw = None
+ layer_idx = 3
+ scaling = HEAD_DIM**-0.5
+
+ def __init__(self, seq, paged_out):
+ self._qkv = (
+ torch.randn(1, N_HEADS, seq, HEAD_DIM),
+ torch.randn(1, KV_HEADS, seq, HEAD_DIM),
+ torch.randn(1, KV_HEADS, seq, HEAD_DIM),
+ )
+ self.paged = _PagedProbe(paged_out)
+ self.sdpa_calls: list[torch.Tensor | None] = []
+
+ def _project_and_rope(self, *args, **kwargs):
+ return self._qkv
+
+ def _run_attn(self, query, key, value, mask):
+ self.sdpa_calls.append(mask)
+ return torch.zeros(1, N_HEADS, query.shape[2], HEAD_DIM)
+
+ def o_proj(self, hidden_states):
+ return hidden_states * 2, None
+
+
+def test_forward_und_sends_a_single_token_decode_to_the_paged_cache():
+ """Goes through the real dispatch, so deleting the branch turns this red."""
+ out = torch.randn(1, N_HEADS, 1, HEAD_DIM)
+ host = _AttnHost(seq=1, paged_out=out)
+ got = host.forward_und(torch.zeros(1, 1, 16), None, None, position_embeddings=None, paged_cache=host.paged)
+ assert len(host.paged.calls) == 1, "the paged hook was never reached"
+ assert host.sdpa_calls == [], "the same step also ran the SDPA path"
+ assert host.paged.calls[0][0] == host.layer_idx
+ torch.testing.assert_close(got, out.reshape(1, 1, -1).contiguous() * 2)
+
+
+@pytest.mark.parametrize(
+ "seq,mask",
+ [(2, None), (1, torch.zeros(1, 1, 1, 1))],
+ ids=["multi_token", "masked"],
+)
+def test_only_an_unmasked_single_token_step_takes_the_paged_path(seq, mask):
+ """Both other shapes are uncapturable, and the cache holds one row per step."""
+ host = _AttnHost(seq=seq, paged_out=torch.zeros(1, N_HEADS, seq, HEAD_DIM))
+ host.forward_und(torch.zeros(1, seq, 16), None, mask, position_embeddings=None, paged_cache=host.paged)
+ assert host.paged.calls == [], "a step that cannot be captured reached the paged cache"
+ assert len(host.sdpa_calls) == 1
+
+
+def test_a_kernel_missing_any_kwarg_we_pass_is_not_supported(monkeypatch):
+ """The call names every argument, so the probe has to cover every name.
+
+ An older wheel exports ``flash_attn_varlen_func`` without the paged kwargs,
+ and a future one may rename the standard ones. Either way the probe is what
+ keeps that install on SDPA instead of raising TypeError at the first decode
+ step, with the feature on by default.
+ """
+
+ def no_paged_kwargs(q, k, v, max_seqlen_q, cu_seqlens_q, max_seqlen_k, cu_seqlens_k, causal=False):
+ raise AssertionError("must not be called") # pragma: no cover
+
+ def renamed_standard_args(q, k, v, max_q, cu_q, max_k, seqused_k=None, block_table=None, causal=False):
+ raise AssertionError("must not be called") # pragma: no cover
+
+ for kernel in (no_paged_kwargs, renamed_standard_args):
+ monkeypatch.setattr(paged_decode, "_flash_varlen", lambda kernel=kernel: kernel)
+ assert paged_decode.paged_decode_supported(torch.device("cuda"), HEAD_DIM) is False
+
+ def current_kernel(
+ q,
+ k,
+ v,
+ max_seqlen_q,
+ cu_seqlens_q,
+ max_seqlen_k,
+ seqused_k=None,
+ block_table=None,
+ softmax_scale=None,
+ causal=False,
+ ):
+ raise AssertionError("must not be called") # pragma: no cover
+
+ monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: current_kernel)
+ assert paged_decode.paged_decode_supported(torch.device("cuda"), HEAD_DIM) is True
+
+
+def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch):
+ """A positional standard argument binds to whatever a future wheel puts in
+ that slot, and the probe above only checks that the names exist."""
+ seen = {}
+
+ def recorder(*args, **kwargs):
+ seen["args"], seen["kwargs"] = args, kwargs
+ return torch.zeros(1, N_HEADS, HEAD_DIM)
+
+ monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: recorder)
+ cache = PagedDecodeCache(1, KV_HEADS, HEAD_DIM, 8, torch.device("cpu"), torch.float32)
+ cache.set_length(4)
+ cache.attend(
+ 0,
+ torch.zeros(1, N_HEADS, 1, HEAD_DIM),
+ torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
+ torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
+ 1.0,
+ )
+ assert len(seen["args"]) == 3, "only q, k and v may be positional"
+ assert set(seen["kwargs"]) == {
+ "max_seqlen_q",
+ "cu_seqlens_q",
+ "max_seqlen_k",
+ "seqused_k",
+ "block_table",
+ "softmax_scale",
+ "causal",
+ }
+
+
+def test_the_kill_switch_returns_no_cache(monkeypatch):
+ """The recipe documents `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` as the way back
+ to the ordinary cache, so the switch needs a test rather than only prose."""
+ from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+ SenseNovaU1Pipeline,
+ )
+
+ monkeypatch.setenv("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "0")
+ # The env check returns before anything else on the pipeline is touched.
+ host = object.__new__(SenseNovaU1Pipeline)
+ assert SenseNovaU1Pipeline._decode_context(host, object()) is None
+
+
+def _pipeline_host(monkeypatch, device="cpu", language_model=None):
+ """A bare pipeline carrying only what ``_decode_context`` touches."""
+ from vllm_omni.diffusion.models.sensenova_u1 import pipeline_sensenova_u1 as pipe_mod
+
+ monkeypatch.setenv("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "1")
+ monkeypatch.setattr(pipe_mod, "paged_decode_supported", lambda dev, head_dim: True)
+ layer = SimpleNamespace(self_attn=SimpleNamespace(head_dim=HEAD_DIM))
+ host = object.__new__(pipe_mod.SenseNovaU1Pipeline)
+ # The pipeline is an nn.Module, and the language model assigned below is one
+ # too, so the bookkeeping has to exist before the assignment.
+ torch.nn.Module.__init__(host)
+ host.device = device
+ lm = language_model if language_model is not None else torch.nn.Module()
+ lm.model = SimpleNamespace(layers=[layer] * LAYERS)
+ host.language_model = lm
+ return pipe_mod, host
+
+
+def test_the_decode_context_is_reused_across_requests(monkeypatch):
+ """A capture binds the addresses it recorded, so a cache built per request
+ forces a capture per request. Two requests whose prefill fits the same
+ bucket have to come back with the same cache and the same runner, and with
+ the generation unchanged so the captures stay valid.
+ """
+ pipe_mod, host = _pipeline_host(monkeypatch)
+ monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
+ ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+ first = ctx(host, _dyn_cache(10))
+ assert first is not None
+ second_input = _dyn_cache(20)
+ second = ctx(host, second_input)
+ assert second[0] is first[0], "the cache was rebuilt for a request that fits it"
+ assert second[1] is first[1], "the runner was rebuilt, so its captures are gone"
+ assert second[0].generation == first[0].generation, "reuse must not invalidate a capture"
+ assert second[0].length == 20
+
+ # The reused buffers must carry the new request's prefill, not the old one.
+ flat = second[0].k[0].view(-1, KV_HEADS, HEAD_DIM)[:20]
+ torch.testing.assert_close(flat, second_input.layers[0].keys[0].transpose(0, 1))
+
+ # A prefill past the bucket cannot share the buffers.
+ bigger = ctx(host, _dyn_cache(BUCKETS[0] + 1))
+ assert bigger[0] is not first[0], "a prefill past the bucket must get new buffers"
+
+
+def test_releasing_the_captures_forces_a_rebuild(monkeypatch):
+ """Sleep level 2 discards the memory a capture recorded, so what the
+ pipeline reuses across requests has to be droppable."""
+ pipe_mod, host = _pipeline_host(monkeypatch)
+ monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
+ ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+ first = ctx(host, _dyn_cache(10))
+ assert ctx(host, _dyn_cache(10))[0] is first[0], "the fixture is not reusing to begin with"
+ pipe_mod.SenseNovaU1Pipeline.release_captured_graphs(host)
+ after = ctx(host, _dyn_cache(10))
+ assert after[0] is not first[0], "the cache survived the release"
+ assert after[1] is not first[1], "the runner survived the release"
+
+
+class _FakeLoRAWrapper(BaseLayerWithLoRA):
+ """Stands in for what ``replace_submodule`` leaves in the module tree."""
+
+ def __init__(self):
+ super().__init__()
+
+
+def test_no_capture_crosses_a_request_once_lora_wrappers_are_installed(monkeypatch):
+ """Base -> LoRA -> base. The reuse test cannot see an adapter change.
+
+ ``DiffusionLoRAManager`` replaces the decode path's linear layers with
+ ``BaseLayerWithLoRA`` and binds, rescales or resets slot 0 per request, and
+ the wrappers stay in the tree afterwards. A graph captured before an adapter
+ was bound would replay without its matmuls; one captured with it bound would
+ survive its removal. Neither moves the shape, dtype or bucket that decide
+ reuse, so the wrappers themselves have to end it.
+ """
+ pipe_mod, host = _pipeline_host(monkeypatch)
+ monkeypatch.setattr(pipe_mod, "DecodeGraphRunner", lambda lm, cache, dev: SimpleNamespace(cache=cache))
+ ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+ base = ctx(host, _dyn_cache(10))
+ assert ctx(host, _dyn_cache(10))[0] is base[0], "the fixture is not reusing to begin with"
+
+ host.language_model.add_module("q_proj", _FakeLoRAWrapper())
+ adapted = ctx(host, _dyn_cache(10))
+ assert adapted[0] is not base[0], "an adapted request reused the base capture"
+ assert adapted[1] is not base[1], "an adapted request reused the base runner"
+
+ back = ctx(host, _dyn_cache(10))
+ assert back[0] is not adapted[0], "a later request reused the adapted capture"
+ assert back[1] is not adapted[1], "a later request reused the adapted runner"
+ assert getattr(host, "_paged_decode", None) is None, "an adapted capture was left on the pipeline"
+
+
+@cuda_only
+@pytest.mark.cuda
+@pytest.mark.L4
+def test_a_second_request_replays_the_first_capture(monkeypatch):
+ """Reuse is only worth anything if the capture survives it."""
+ dev = torch.device("cuda")
+ pipe_mod, host = _pipeline_host(monkeypatch, device="cuda", language_model=_StubLM(dev))
+ ctx = pipe_mod.SenseNovaU1Pipeline._decode_context
+
+ cache, runner = ctx(host, _dyn_cache(10, dev))
+ cache.set_length(cache.length + 1)
+ runner.step(0, 0)
+ assert runner.captures == 1
+
+ cache2, runner2 = ctx(host, _dyn_cache(12, dev))
+ assert cache2 is cache and runner2 is runner
+ cache2.set_length(cache2.length + 1)
+ runner2.step(0, 0)
+ assert runner.captures == 1, "the second request captured the graph again"
+
+
+class _Tok:
+ """The two tokenizer calls ``_generate_text`` makes."""
+
+ EOS = 7
+
+ def convert_tokens_to_ids(self, token):
+ return self.EOS
+
+ def decode(self, ids, skip_special_tokens=False):
+ return " ".join(str(i) for i in ids)
+
+
+class _Out:
+ def __init__(self, logits, past_key_values):
+ self.logits = logits
+ self.past_key_values = past_key_values
+
+
+class _TextHost:
+ """Carries only what ``_generate_text`` touches."""
+
+ def __init__(self, context):
+ self._context = context
+ self.seen: list[object] = []
+ self.tokenizer = _Tok()
+
+ def _decode_context(self, past_key_values):
+ return self._context
+
+ def _ar_step(self, next_token, t_idx, past_key_values, decode=None):
+ self.seen.append(decode)
+ # Stop the loop: argmax of this row is the EOS id.
+ logits = torch.full((1, 1, 8), -1.0)
+ logits[0, 0, _Tok.EOS] = 1.0
+ return _Out(logits, past_key_values)
+
+
+def test_text_decoding_uses_the_paged_context_too():
+ """Image-to-text and text-to-text run their own decode loop. It used to call
+ ``_ar_step`` without the context, so the paged path was documented but never
+ reached outside think."""
+ from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+ SenseNovaU1Pipeline,
+ )
+
+ context = object()
+ host = _TextHost(context)
+ prefix = torch.full((1, 1, 8), -1.0)
+ prefix[0, 0, 5] = 1.0 # not EOS, so one step runs
+ SenseNovaU1Pipeline._generate_text(host, prefix, past_key_values=None, t_idx=0)
+ assert host.seen == [context], f"text decoding ran with decode={host.seen}"
+
+
+class _WarmupReq:
+ """The engine's dummy request, which carries the reserved id."""
+
+ def __init__(self, request_id):
+ self.request_id = request_id
+ self.prompts = [{"prompt": "", "modalities": ["image"]}]
+
+
+def test_the_dummy_warmup_request_drives_one_decode_step(monkeypatch):
+ """The dummy request is a text-to-image run with think off, so it exercises
+ the prefill shape and never the decode one. Without this hook whatever the
+ deploy config compiles for decode lands on the first real request."""
+ from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+ SenseNovaU1Pipeline,
+ )
+ from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID
+
+ calls: list[str] = []
+ host = object.__new__(SenseNovaU1Pipeline)
+ monkeypatch.setattr(host, "_warm_ar_decode", lambda: calls.append("warm"), raising=False)
+ monkeypatch.setattr(
+ SenseNovaU1Pipeline,
+ "_parse_request",
+ lambda self, req: (_ for _ in ()).throw(_StopForwardError()),
+ raising=False,
+ )
+
+ with pytest.raises(_StopForwardError):
+ SenseNovaU1Pipeline.forward(host, _WarmupReq(DUMMY_DIFFUSION_REQUEST_ID))
+ assert calls == ["warm"], "the dummy warmup request did not drive a decode step"
+
+ calls.clear()
+ with pytest.raises(_StopForwardError):
+ SenseNovaU1Pipeline.forward(host, _WarmupReq("a-real-request"))
+ assert calls == [], "a real request paid for the warmup"
+
+
+class _StopForwardError(Exception):
+ """Stops ``forward`` right after the warmup hook, so the test needs no model."""
diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_rmsnorm.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_rmsnorm.py
new file mode 100644
index 000000000..ddb4bd61e
--- /dev/null
+++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_rmsnorm.py
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for SenseNova-U1 RMSNorm and the U1.5 config surface.
+
+``Qwen3RMSNorm`` was a hand-rolled fp32 up/down-cast chain; it now calls the fused
+``F.rms_norm``. These pin the mathematics against an independent reference so a
+future kernel swap cannot silently change what the layer computes, and pin the two
+config fields that distinguish U1.5 from U1.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import Qwen3RMSNorm
+from vllm_omni.transformers_utils.configs.sensenova_u1 import SenseNovaU1Config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _reference(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+ """RMSNorm written out in float64 so it is independent of the implementation."""
+ x64 = x.to(torch.float64)
+ var = x64.pow(2).mean(-1, keepdim=True)
+ return (weight.to(torch.float64) * (x64 * torch.rsqrt(var + eps))).to(x.dtype)
+
+
+@pytest.mark.parametrize("hidden", [8, 128, 4096])
+@pytest.mark.parametrize(
+ "dtype,atol",
+ [(torch.float32, 1e-5), (torch.bfloat16, 3e-2), (torch.float16, 4e-3)],
+)
+def test_rmsnorm_matches_float64_reference(hidden: int, dtype: torch.dtype, atol: float):
+ torch.manual_seed(0)
+ norm = Qwen3RMSNorm(hidden, eps=1e-6).to(dtype)
+ with torch.no_grad():
+ norm.weight.copy_(torch.rand(hidden, dtype=dtype) + 0.5)
+ x = torch.randn(7, hidden, dtype=dtype)
+ got = norm(x)
+ want = _reference(x, norm.weight, norm.variance_epsilon)
+ assert got.dtype == dtype
+ assert got.shape == x.shape
+ torch.testing.assert_close(got.to(torch.float32), want.to(torch.float32), atol=atol, rtol=atol)
+
+
+def _cast_chain(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+ """The implementation this PR replaced, kept so its accuracy can be compared."""
+ var = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+ return weight * (x.to(torch.float32) * torch.rsqrt(var + eps)).to(weight.dtype)
+
+
+@pytest.mark.parametrize("hidden", [512, 4096])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_rmsnorm_is_closer_to_float64_than_the_cast_chain(hidden: int, dtype: torch.dtype):
+ """The cast chain rounds to ``dtype`` before the weight multiply and so rounds
+ twice, while ``F.rms_norm`` rounds once. Restoring the cast chain fails here."""
+ torch.manual_seed(0)
+ norm = Qwen3RMSNorm(hidden, eps=1e-6).to(dtype)
+ with torch.no_grad():
+ norm.weight.copy_(torch.rand(hidden, dtype=dtype) + 0.5)
+ x = torch.randn(64, hidden, dtype=dtype)
+
+ reference = _reference(x, norm.weight, norm.variance_epsilon).to(torch.float64)
+ scale = reference.abs().mean()
+
+ def relative_error(out: torch.Tensor) -> float:
+ return ((out.to(torch.float64) - reference).abs().mean() / scale).item()
+
+ fused = relative_error(norm(x))
+ legacy = relative_error(_cast_chain(x, norm.weight, norm.variance_epsilon))
+ assert fused < legacy, f"{dtype} hidden={hidden}: fused {fused:.3e} not better than {legacy:.3e}"
+
+
+def test_rmsnorm_is_scale_equivariant():
+ """RMSNorm divides by the RMS, so scaling the input must not change the output."""
+ torch.manual_seed(0)
+ norm = Qwen3RMSNorm(64, eps=1e-12)
+ x = torch.randn(4, 64)
+ torch.testing.assert_close(norm(x), norm(x * 7.5), atol=1e-4, rtol=1e-4)
+
+
+def test_rmsnorm_preserves_leading_dims():
+ norm = Qwen3RMSNorm(16, eps=1e-6)
+ for shape in [(16,), (3, 16), (2, 3, 16)]:
+ assert norm(torch.randn(*shape)).shape == shape
+
+
+def test_rmsnorm_weight_is_applied():
+ norm = Qwen3RMSNorm(4, eps=1e-6)
+ x = torch.randn(2, 4)
+ base = norm(x)
+ with torch.no_grad():
+ norm.weight.mul_(3.0)
+ torch.testing.assert_close(norm(x), base * 3.0, atol=1e-5, rtol=1e-5)
+
+
+def test_u15_config_fields_default_to_u1_values():
+ """U1.5 flips exactly these two; the defaults must stay U1's so old checkpoints
+ keep their behaviour when the keys are absent."""
+ cfg = SenseNovaU1Config()
+ assert cfg.use_pixel_head is False
+ assert cfg.noise_scale_max_value == 10.0
+
+
+def test_u15_config_fields_are_read_from_kwargs():
+ cfg = SenseNovaU1Config(use_pixel_head=True, noise_scale_max_value=16.0)
+ assert cfg.use_pixel_head is True
+ assert cfg.noise_scale_max_value == 16.0
diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py
index 814450d92..2626b1258 100644
--- a/tests/diffusion/test_diffusion_model_runner.py
+++ b/tests/diffusion/test_diffusion_model_runner.py
@@ -179,6 +179,30 @@ def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
assert events == [backend]
+def test_release_captured_graphs_clears_runners_and_delegates_to_the_pipeline():
+ """Both halves in one place, so a pipeline that keeps captures is collected
+ without the caller knowing which pipelines have one."""
+ released = []
+ runner = object.__new__(DiffusionModelRunner)
+ runner.graph_runners = {"decode": object()}
+ runner.pipeline = SimpleNamespace(release_captured_graphs=lambda: released.append(True))
+
+ runner.release_captured_graphs()
+
+ assert runner.graph_runners == {}
+ assert released == [True]
+
+
+def test_release_captured_graphs_tolerates_a_pipeline_without_captures():
+ """Most pipelines keep none, and the runner may not carry `graph_runners`."""
+ runner = object.__new__(DiffusionModelRunner)
+ runner.pipeline = object()
+
+ runner.release_captured_graphs()
+
+ assert not hasattr(runner, "graph_runners")
+
+
def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py
index 0d8d5e92a..52717fb6f 100644
--- a/tests/diffusion/test_diffusion_worker.py
+++ b/tests/diffusion/test_diffusion_worker.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for DiffusionWorker class.
@@ -143,6 +143,21 @@ class TestDiffusionWorkerSleep:
assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
+ def test_sleep_level_2_releases_captures_through_the_model_runner(self, mocker: MockerFixture, mock_gpu_worker):
+ """Level 2 discards the memory captures were recorded against, and the
+ runner owns the release. The worker asks it once and does not reach past
+ it into the pipeline."""
+ mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+ mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
+ mock_platform.get_device_total_memory.return_value = 80 * 1024**3
+ mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory", side_effect=[0, 0])
+ mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(return_value=[])
+
+ mock_gpu_worker.sleep(level=2)
+
+ mock_gpu_worker.model_runner.release_captured_graphs.assert_called_once_with()
+ mock_gpu_worker.model_runner.pipeline.release_captured_graphs.assert_not_called()
+
def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
"""Test that sleep validates memory was actually freed."""
mock_allocator_class = patch_cumem_allocator(mocker)
diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py
index 2a7385c1d..76d1afec6 100644
--- a/vllm_omni/config/environment_variable_inventory.py
+++ b/vllm_omni/config/environment_variable_inventory.py
@@ -167,6 +167,7 @@ _MODEL_PROMOTE = (
"VLLM_GEPARD_LOOKBACK_FRAMES",
"VLLM_OMNI_FISH_KVCACHE_ATTN",
"VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
+ "VLLM_OMNI_SENSENOVA_PAGED_DECODE",
)
_MODEL_REQUEST_SCOPE = (
diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py
index 9a9fbfae5..676e8d514 100644
--- a/vllm_omni/diffusion/lora/loader.py
+++ b/vllm_omni/diffusion/lora/loader.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import os
from collections import OrderedDict, defaultdict
from collections.abc import Callable
@@ -41,7 +44,7 @@ def _prepare_lora_delta(
lora_b_suffix: str = "lora_B.weight",
lora_bias_suffix: str = "bias",
):
- used_keys = set()
+ used_keys: set[str] = set()
# stacked_params_mapping param_to_weight_names
# [(".to_qkv", ".to_q.", "q")
# (".to_qkv", ".to_k.", "k") ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
@@ -54,7 +57,10 @@ def _prepare_lora_delta(
if param_to_weight_names is None:
param_to_weight_names = defaultdict(list)
for param_name, weight_names in param_to_weight_names.items():
- if param_name not in base_key:
+ # Fused names can overlap by prefix (`.qkv_proj` is a substring of
+ # `.qkv_proj_mot_gen`), and a substring match stacks both mappings into
+ # one delta of twice the height.
+ if not base_key.endswith(param_name):
continue
is_stacked_param = True
# handle lora_a_key and lora_b_key together
@@ -78,7 +84,7 @@ def _prepare_lora_delta(
stacked_deltas.append(delta)
used_keys.add(lora_a_key)
used_keys.add(lora_b_key)
- continue
+ break
if is_stacked_param:
return torch.concat(stacked_deltas), used_keys
@@ -200,6 +206,10 @@ def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) ->
class LoraLoaderMixin:
transformer_name = "transformer"
+ # Set by the pipeline this mixin is combined with; annotated, not assigned.
+ transformer: torch.nn.Module
+ _lora_loaded: dict[str | None, dict[str, torch.Tensor]]
+
# Lazy initialization to avoid MRO issues: __init__ may not be called
# when mixin with nn.Module
@property
@@ -373,6 +383,7 @@ class QwenImageLoraLoaderMixin(LoraLoaderMixin):
class WanLoraLoaderMixin(LoraLoaderMixin):
transformer_2_name = "transformer_2"
+ has_transformer_2: bool
def load_lora_weights(
self,
diff --git a/vllm_omni/diffusion/models/sensenova_u1/paged_decode.py b/vllm_omni/diffusion/models/sensenova_u1/paged_decode.py
new file mode 100644
index 000000000..bade5cbe5
--- /dev/null
+++ b/vllm_omni/diffusion/models/sensenova_u1/paged_decode.py
@@ -0,0 +1,368 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Paged KV cache for SenseNova-U1 autoregressive decode, and a CUDA graph over it.
+
+Autoregressive decode dominates a think request -- 80% of the wall clock, at
+about 20 ms per token for 42 layers -- and almost all of the GPU idle in that
+window is Python/ATen dispatch rather than any CUDA call. A CUDA graph retires
+those dispatches, but capture needs static shapes, and the default cache grows
+K/V with ``torch.cat`` on every step.
+
+Padding the cache to a bucket and masking the tail does give static shapes, but
+the mask is what costs: measured on an A800 at real decode shapes, a masked
+bucket attention runs 7.96-11.55 ms per step against 1.03 ms for an exact-length
+unmasked one, which is more than the graph saves.
+
+The way out is a paged cache. ``flash_attn_varlen_func`` takes the used length
+as a *tensor* (``seqused_k``) alongside a ``block_table``, so the buffers stay
+bucket-sized and capturable while the kernel reads only the valid prefix. One
+captured graph then serves every length in the bucket.
+
+Scope, and when to delete this. The cache is model-local on purpose:
+``DiffusionKVCacheManager`` reserves once per scheduler request, and
+``ARDiffusionModelRunner`` lives under ``vllm_omni/experimental``. SenseNova runs
+its whole autoregressive loop inside one pipeline ``forward``, so the decode
+steps never surface to the scheduler and it cannot admit, pool, evict, reuse a
+prefix for, or continuously batch them. What is here is therefore one set of
+buffers behind an identity block table, reused by whichever request fits them,
+not general paged-KV support: it holds while the pipeline serves one sequence
+per forward, and is released when sleep discards the memory it captured
+against. Once the decode loop is
+driven by a scheduler-visible runner, or the pipeline handles more than one
+sequence per forward, delete this file and use the manager instead.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+BLOCK_SIZE = 16
+# Growth schedule. Each new bucket costs one capture, so keep the list short and
+# let the last entry cover the tail rather than doubling forever.
+BUCKETS = (512, 1024, 2048, 4096, 8192)
+# Past the last bucket the schedule grows by this much at a time. A request
+# decodes at most `max_think_tokens` (1024) or `max_tokens` (512) steps, so a
+# step this size costs at most one reallocation per request. Doubling the last
+# bucket instead would also cost at most one, but a 9,100-token request would
+# then reserve 766 MiB it never uses, against 118 MiB here.
+TAIL_STEP = 2048
+
+
+# Every argument to the kernel is passed by keyword, so the probe covers all of
+# them. An older wheel can export `flash_attn_varlen_func` without the paged
+# ones, or under different names, and a version string would not say which;
+# probe the signature and fall back to SDPA when any is missing.
+_REQUIRED_KWARGS = (
+ "max_seqlen_q",
+ "cu_seqlens_q",
+ "max_seqlen_k",
+ "seqused_k",
+ "block_table",
+ "softmax_scale",
+ "causal",
+)
+
+
+def _flash_varlen():
+ """Import lazily: the bundled kernel is CUDA-only and optional."""
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+ return flash_attn_varlen_func
+
+
+def _accepts_paged_kwargs(fn) -> bool:
+ try:
+ params = inspect.signature(fn).parameters
+ except (TypeError, ValueError): # pragma: no cover - no introspectable signature
+ return False
+ return all(name in params for name in _REQUIRED_KWARGS)
+
+
+def paged_decode_supported(device: torch.device, head_dim: int) -> bool:
+ if device.type != "cuda":
+ return False
+ if head_dim % 8 or head_dim > 256:
+ return False
+ try:
+ fn = _flash_varlen()
+ except Exception as exc: # pragma: no cover - depends on the wheel
+ logger.debug("Paged decode unavailable: %s", exc)
+ return False
+ if not _accepts_paged_kwargs(fn):
+ logger.debug("Paged decode unavailable: flash_attn_varlen_func lacks %s", _REQUIRED_KWARGS)
+ return False
+ return True
+
+
+def dynamic_lora_wrappers_present(module) -> bool:
+ """True once ``DiffusionLoRAManager`` has wrapped layers under ``module``.
+
+ The manager replaces linear layers with ``BaseLayerWithLoRA`` and then
+ binds, rescales or resets slot 0 per request; the wrappers stay in the tree
+ afterwards. Capture records the module tree and the branches that ran during
+ it, so a graph taken before an adapter was bound would replay without the
+ adapter's matmuls and one taken with it bound would survive its removal --
+ both silently, and neither visible in the shape, dtype and bucket that decide
+ reuse. Nothing captured may outlive a request once these exist.
+
+ The distilled LoRA is not affected: it is fused one-way into the weights and
+ leaves no wrapper behind.
+ """
+ try:
+ from vllm.lora.layers import BaseLayerWithLoRA
+ except ImportError: # pragma: no cover - depends on the wheel
+ return False
+ return any(isinstance(m, BaseLayerWithLoRA) for m in module.modules())
+
+
+def _bucket_for(length: int) -> int:
+ for b in BUCKETS:
+ if length <= b:
+ return b
+ # Past the schedule, grow in TAIL_STEP-sized steps. Rounding to a whole
+ # number of blocks here would reallocate and re-capture every BLOCK_SIZE
+ # tokens: a think edit with two 2048x2048 inputs runs to about 9,100 tokens
+ # and paid 41 reallocations and 42 captures in one request.
+ over = length - BUCKETS[-1]
+ return BUCKETS[-1] + ((over + TAIL_STEP - 1) // TAIL_STEP) * TAIL_STEP
+
+
+class PagedDecodeCache:
+ """Per-layer paged K/V for single-token decode.
+
+ Layout follows what ``flash_attn_varlen_func`` accepts with a block table:
+ ``(num_blocks, BLOCK_SIZE, kv_heads, head_dim)``. ``seqused`` holds the
+ number of valid tokens and is the only thing that changes between decode
+ steps, which is what lets a captured graph be replayed unchanged.
+ """
+
+ def __init__(self, num_layers, kv_heads, head_dim, length, device, dtype):
+ self.kv_heads = kv_heads
+ self.head_dim = head_dim
+ self.device = device
+ self.dtype = dtype
+ self.bucket = _bucket_for(length)
+ nblocks = self.bucket // BLOCK_SIZE
+ shape = (nblocks, BLOCK_SIZE, kv_heads, head_dim)
+ self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(num_layers)]
+ self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(num_layers)]
+ self.block_table = torch.arange(nblocks, device=device, dtype=torch.int32).unsqueeze(0)
+ self.cu_seqlens_q = torch.tensor([0, 1], device=device, dtype=torch.int32)
+ self.seqused = torch.zeros(1, device=device, dtype=torch.int32)
+ # The write slot has to live in a tensor, not a Python int. Under graph
+ # capture a Python index is baked in at capture time, so every replay
+ # would overwrite the same slot -- correct eagerly, silently wrong once
+ # captured.
+ self.pos = torch.zeros(1, device=device, dtype=torch.int64)
+ self._length = 0
+ # Bumped whenever the buffers are reallocated, so a graph runner can
+ # tell that its capture no longer points at live memory.
+ self.generation = 0
+
+ @property
+ def length(self) -> int:
+ return self._length
+
+ def set_length(self, n: int) -> None:
+ self._length = n
+ self.seqused.fill_(n)
+ self.pos.fill_(n - 1)
+
+ def reusable_for(self, num_layers, kv_heads, head_dim, dtype, length) -> bool:
+ """Can this cache serve another request without reallocating?
+
+ Reusing the buffers is what lets captured graphs survive across
+ requests: a capture binds the addresses it recorded, so a fresh cache
+ forces a fresh capture.
+ """
+ return (
+ len(self.k) == num_layers
+ and self.kv_heads == kv_heads
+ and self.head_dim == head_dim
+ and self.dtype == dtype
+ and _bucket_for(length) <= self.bucket
+ )
+
+ def load_prefix(self, cache) -> bool:
+ """Copy a prefill cache into the paged buffers.
+
+ Only the first ``seqused`` rows are ever read, so the tail is left as
+ whatever the previous request wrote. Returns False for an empty cache.
+ """
+ layer0 = cache.layers[0]
+ if layer0.keys is None:
+ return False
+ prefix = layer0.keys.shape[2]
+ for i in range(len(self.k)):
+ keys, values = cache.layers[i].keys, cache.layers[i].values
+ # [B, H, S, D] -> [S, H, D], then into the paged blocks
+ self._write_prefix(i, keys[0].transpose(0, 1), values[0].transpose(0, 1))
+ self.set_length(prefix)
+ return True
+
+ @classmethod
+ def from_dynamic_cache(cls, cache, num_layers, device, dtype, min_length=0):
+ """Allocate buffers for a prefill cache and copy it in."""
+ layer0 = cache.layers[0]
+ if layer0.keys is None:
+ return None
+ _, kv_heads, prefix, head_dim = layer0.keys.shape
+ obj = cls(num_layers, kv_heads, head_dim, max(prefix + 1, min_length), device, dtype)
+ obj.load_prefix(cache)
+ return obj
+
+ def _write_prefix(self, layer_idx, keys_shd, values_shd):
+ n = keys_shd.shape[0]
+ flat_k = self.k[layer_idx].view(-1, self.kv_heads, self.head_dim)
+ flat_v = self.v[layer_idx].view(-1, self.kv_heads, self.head_dim)
+ flat_k[:n].copy_(keys_shd)
+ flat_v[:n].copy_(values_shd)
+
+ def grow(self, length: int) -> bool:
+ """Move to the next bucket, keeping the tokens already stored.
+
+ Everything is built into locals first and published in one step. A
+ resize that failed part way -- an allocation for a later layer, say --
+ would otherwise leave new K beside old V under the ``generation`` a
+ captured graph is keyed on, so the next request would keep selecting
+ that graph and replay it over storage that had just been freed.
+ """
+ new_bucket = _bucket_for(length)
+ if new_bucket == self.bucket:
+ return False
+ nblocks = new_bucket // BLOCK_SIZE
+ shape = (nblocks, BLOCK_SIZE, self.kv_heads, self.head_dim)
+ old_len = self._length
+ new_k, new_v = [], []
+ for old_k, old_v in zip(self.k, self.v):
+ k = torch.zeros(shape, device=self.device, dtype=self.dtype)
+ v = torch.zeros(shape, device=self.device, dtype=self.dtype)
+ k.view(-1, self.kv_heads, self.head_dim)[:old_len].copy_(
+ old_k.view(-1, self.kv_heads, self.head_dim)[:old_len]
+ )
+ v.view(-1, self.kv_heads, self.head_dim)[:old_len].copy_(
+ old_v.view(-1, self.kv_heads, self.head_dim)[:old_len]
+ )
+ new_k.append(k)
+ new_v.append(v)
+ block_table = torch.arange(nblocks, device=self.device, dtype=torch.int32).unsqueeze(0)
+
+ # Publish. New buffers mean any graph captured against the old ones is
+ # stale, so the generation moves in the same step that hands them over.
+ self.k, self.v, self.block_table, self.bucket = new_k, new_v, block_table, new_bucket
+ self.generation += 1
+ self.set_length(old_len)
+ logger.debug("Paged decode grew to bucket %d at length %d (generation %d)", new_bucket, length, self.generation)
+ return True
+
+ def attend(self, layer_idx, query_bhsd, key_bhsd, value_bhsd, softmax_scale):
+ """Append this step's K/V and attend over the whole valid prefix.
+
+ Writes at ``length - 1`` because the caller has already advanced the
+ length for this step. Both the slot and the attended length come from
+ device tensors, so a captured graph reads whatever they hold at replay
+ rather than whatever they held at capture.
+ """
+ flat_k = self.k[layer_idx].view(-1, self.kv_heads, self.head_dim)
+ flat_v = self.v[layer_idx].view(-1, self.kv_heads, self.head_dim)
+ flat_k.index_copy_(0, self.pos, key_bhsd[0, :, 0].unsqueeze(0))
+ flat_v.index_copy_(0, self.pos, value_bhsd[0, :, 0].unsqueeze(0))
+ heads, head_dim = query_bhsd.shape[1], query_bhsd.shape[3]
+ q = query_bhsd[0].transpose(0, 1).reshape(1, heads, head_dim)
+ out = _flash_varlen()(
+ q,
+ self.k[layer_idx],
+ self.v[layer_idx],
+ max_seqlen_q=1,
+ cu_seqlens_q=self.cu_seqlens_q,
+ max_seqlen_k=self.bucket,
+ seqused_k=self.seqused,
+ block_table=self.block_table,
+ softmax_scale=softmax_scale,
+ causal=False,
+ )
+ return out.reshape(1, 1, heads, head_dim)
+
+ def to_dynamic_cache(self, cache) -> None:
+ """Write the decoded K/V back so the generation stage sees one cache.
+
+ Only decode is paged; the DiT stage that follows reads the ordinary
+ cache, so the two are reconciled once at the hand-off rather than kept
+ in sync on every step.
+ """
+ n = self._length
+ for i, layer in enumerate(cache.layers):
+ keys = self.k[i].view(-1, self.kv_heads, self.head_dim)[:n]
+ values = self.v[i].view(-1, self.kv_heads, self.head_dim)[:n]
+ layer.keys = keys.transpose(0, 1).unsqueeze(0).contiguous()
+ layer.values = values.transpose(0, 1).unsqueeze(0).contiguous()
+
+
+class DecodeGraphRunner:
+ """One captured decode step per KV bucket, replayed for every token.
+
+ The whole point of the paged cache is that everything the step reads which
+ varies -- the token, its position, and how much of the cache is live -- sits
+ in device tensors. So a single capture serves every step in a bucket: fill
+ the tensors, replay, read the logits out of the static output.
+
+ A capture is invalidated when the cache reallocates (tracked by
+ ``PagedDecodeCache.generation``), which happens once per bucket boundary.
+ """
+
+ def __init__(self, language_model, cache, device):
+ self.lm = language_model
+ self.cache = cache
+ self.device = device
+ self.input_ids = torch.zeros(1, 1, dtype=torch.long, device=device)
+ self.indexes = torch.zeros(3, 1, dtype=torch.long, device=device)
+ self._graphs: dict[tuple[int, int], tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {}
+ self.captures = 0
+
+ def _forward(self):
+ return self.lm(
+ input_ids=self.input_ids,
+ indexes=self.indexes,
+ past_key_values=None,
+ use_cache=False,
+ paged_cache=self.cache,
+ ).logits
+
+ def _capture(self, key):
+ # Warm up on a side stream first so cuBLAS workspaces and any lazy
+ # allocation happen outside the capture. The warm-up recomputes this
+ # same step, writing the same K/V to the same slot, so it is idempotent.
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ for _ in range(3):
+ self._forward()
+ torch.cuda.current_stream().wait_stream(side)
+
+ graph = torch.cuda.CUDAGraph()
+ # A private pool would hand every runner its own arena and never give it
+ # back, and a runner is rebuilt whenever the cache reallocates or a LoRA
+ # wrapper appears: measured, that grew device memory by about 40 MiB per
+ # request and never levelled off. The platform pool is shared, which is
+ # what every other captured path in the tree uses.
+ with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+ logits = self._forward()
+ self.captures += 1
+ logger.debug("Captured decode graph for bucket=%d generation=%d", self.cache.bucket, self.cache.generation)
+ self._graphs[key] = (graph, logits)
+ return self._graphs[key]
+
+ def step(self, token, t_index):
+ """Run one decode step. Returns the static logits tensor."""
+ self.input_ids[0, 0] = token
+ self.indexes[0, 0] = t_index
+ key = (self.cache.bucket, self.cache.generation)
+ entry = self._graphs.get(key) or self._capture(key)
+ entry[0].replay()
+ return entry[1]
diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
index 3a828fa69..de636fb6d 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""SenseNova-U1 Pipeline for vLLM-Omni.
SenseNova-U1 is a unified Qwen3-based model that uses Mixture-of-Tokenizers
@@ -35,6 +35,13 @@ from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.lora.loader import (
+ LoraLoaderMixin,
+ _apply_diffusers_lora_alpha_scaling,
+ _load_lora_state_dict,
+ _prepare_lora_delta,
+ _remap_state_dict_keys,
+)
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -44,7 +51,14 @@ from vllm_omni.transformers_utils.configs.sensenova_u1 import (
SenseNovaU1Config,
)
+from .paged_decode import (
+ DecodeGraphRunner,
+ PagedDecodeCache,
+ dynamic_lora_wrappers_present,
+ paged_decode_supported,
+)
from .sensenova_u1_transformer import (
+ SenseNovaU1CausalLMOutput,
SenseNovaU1ForCausalLM,
clear_flash_kv_cache,
create_block_causal_mask,
@@ -446,7 +460,13 @@ class SenseNovaU1DenoisingAdapter(nn.Module):
return self.language_model(*args, **kwargs)
-class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin):
+class SenseNovaU1Pipeline(
+ nn.Module,
+ SupportsComponentDiscovery,
+ DiffusionPipelineProfilerMixin,
+ CFGParallelMixin,
+ LoraLoaderMixin,
+):
"""SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni.
Builds the full model graph internally:
@@ -709,8 +729,71 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
sigma = shift * sigma / (1 + (shift - 1) * sigma)
return 1 - sigma
- def _ar_step(self, next_token, t_idx, past_key_values):
+ def _decode_context(self, past_key_values):
+ """Paged cache plus a captured graph for the AR loop, or ``None``.
+
+ Decode is where a think request spends most of its time, and every step
+ issues the same work at the same shapes, so it is worth capturing. The
+ cache has to be paged for that: it keeps the buffers a fixed size while
+ the kernel reads only ``seqused_k`` of them, which is what lets one
+ capture serve every step until the next bucket.
+
+ Both live on the pipeline and are reused across requests unless a LoRA
+ adapter is being served dynamically. A capture binds the addresses it
+ recorded, so a cache built per request forces a capture per request:
+ measured, that left about 40 MiB of device memory behind on every
+ request and re-captured every time. Reuse rests on the same thing the
+ paged path already rests on -- one sequence in flight per pipeline
+ forward.
+ """
+ if os.environ.get("VLLM_OMNI_SENSENOVA_PAGED_DECODE", "1") != "1":
+ return None
+ head_dim = self.language_model.model.layers[0].self_attn.head_dim
+ if not paged_decode_supported(torch.device(self.device), head_dim):
+ return None
+ if past_key_values is None or not past_key_values.layers:
+ return None
+ layer0 = past_key_values.layers[0]
+ if layer0.keys is None:
+ return None
+ _, kv_heads, prefix, kv_head_dim = layer0.keys.shape
+ device = torch.device(self.device)
+ num_layers = len(self.language_model.model.layers)
+ # Reuse is off once the LoRA manager holds the decode path: the test
+ # below sees shape, dtype and bucket, none of which move when an adapter
+ # is bound, rescaled or reset between requests.
+ reuse = not dynamic_lora_wrappers_present(self.language_model)
+ decode = getattr(self, "_paged_decode", None) if reuse else None
+ if decode is None or not decode[0].reusable_for(
+ num_layers, kv_heads, kv_head_dim, layer0.keys.dtype, prefix + 1
+ ):
+ cache = PagedDecodeCache(num_layers, kv_heads, kv_head_dim, prefix + 1, device, layer0.keys.dtype)
+ decode = (cache, DecodeGraphRunner(self.language_model, cache, device))
+ self._paged_decode = decode if reuse else None
+ if not decode[0].load_prefix(past_key_values):
+ return None
+ return decode
+
+ def release_captured_graphs(self) -> None:
+ """Drop the reused paged cache and the graphs captured against it.
+
+ Sleep level 2 discards the memory a capture recorded, so anything held
+ across requests has to go with it. The next request rebuilds both.
+ """
+ self._paged_decode = None
+
+ def _ar_step(self, next_token, t_idx, past_key_values, decode=None):
"""One autoregressive token step. Constructs single-token `indexes` explicitly."""
+ if decode is not None:
+ cache, runner = decode
+ if cache.length + 1 > cache.bucket:
+ cache.grow(cache.length + 1)
+ cache.set_length(cache.length + 1)
+ logits = runner.step(int(next_token), t_idx + 1)
+ # The caller reassigns its cache handle from this field; the paged
+ # tokens are merged into it once, at the hand-off below.
+ return SenseNovaU1CausalLMOutput(logits=logits, past_key_values=past_key_values)
+
indexes = torch.tensor([[t_idx + 1], [0], [0]], dtype=torch.long, device=self.device)
outputs = self.language_model(
input_ids=next_token.unsqueeze(0),
@@ -725,24 +808,28 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
think_end_token_id = self.tokenizer.convert_tokens_to_ids("")
think_token_ids = []
next_token = torch.argmax(prefix_outputs.logits[:, -1, :], dim=-1)
+ decode = self._decode_context(past_key_values)
for _ in range(max_think_tokens):
token_item = next_token.item()
if token_item == eos_token_id:
break
if token_item == think_end_token_id:
- outputs = self._ar_step(next_token, t_idx, past_key_values)
+ outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
past_key_values = outputs.past_key_values
t_idx += 1
think_token_ids.append(token_item)
break
think_token_ids.append(token_item)
- outputs = self._ar_step(next_token, t_idx, past_key_values)
+ outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
past_key_values = outputs.past_key_values
t_idx += 1
next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1)
+ if decode is not None:
+ decode[0].to_dynamic_cache(past_key_values)
+
# Append "\n\n![]()
" tokens to cache
append_ids = self.tokenizer(
"\n\n![]()
",
@@ -797,6 +884,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
"""Autoregressive text decoding seeded from prefix logits."""
eos_token_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
generated_ids: list[int] = []
+ # Same paged decode as the think loop. No write-back at the end: the
+ # caller returns the decoded string and never reads the cache again,
+ # unlike the think path, which hands it to the generation stage.
+ decode = self._decode_context(past_key_values)
logits = prefix_logits[:, -1, :]
for _ in range(max_tokens):
@@ -810,7 +901,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
if token_id == eos_token_id:
break
generated_ids.append(token_id)
- outputs = self._ar_step(next_token, t_idx, past_key_values)
+ outputs = self._ar_step(next_token, t_idx, past_key_values, decode)
past_key_values = outputs.past_key_values
logits = outputs.logits[:, -1, :]
t_idx += 1
@@ -1174,8 +1265,37 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
layer.values = layer.values.expand(batch_size, *layer.values.shape[1:])
prepare_flash_kv_cache(kv, current_len=token_hw, batch_size=batch_size)
+ @staticmethod
+ def _is_warmup_request(req) -> bool:
+ is_dummy_run = getattr(req, "is_dummy_run", None)
+ if callable(is_dummy_run):
+ return bool(is_dummy_run())
+ return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
+
+ def _warm_ar_decode(self) -> None:
+ """Run one single-token decode at startup so its compiled region is built.
+
+ The engine's dummy request is a text-to-image run with think off. It
+ exercises the prefill shape but never the decode one, so whatever the
+ deploy config compiles for decode lands on the first real request
+ instead of on readiness. Warmup is best effort and must not fail startup.
+ """
+ try:
+ lm = self.language_model
+ device = torch.device(self.device)
+ ids = torch.zeros(1, 2, dtype=torch.long, device=device)
+ idx = torch.zeros(3, 2, dtype=torch.long, device=device)
+ prefill = lm(input_ids=ids, indexes=idx, use_cache=True)
+ nxt = torch.zeros(1, 1, dtype=torch.long, device=device)
+ nidx = torch.tensor([[2], [0], [0]], dtype=torch.long, device=device)
+ lm(input_ids=nxt, indexes=nidx, past_key_values=prefill.past_key_values, use_cache=True)
+ except Exception as exc: # pragma: no cover - warmup is best effort
+ logger.warning("Autoregressive decode warmup skipped: %s", exc)
+
@torch.inference_mode()
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+ if self._is_warmup_request(req):
+ self._warm_ar_decode()
p = self._parse_request(req)
input_images = self._extract_input_images(p.first_prompt)
@@ -1408,21 +1528,129 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
# Weight loading
# -----------------------------------------------------------------------
+ # Shard name -> fused parameter, used by weight loading and LoRA loading.
+ # More specific _mot_gen patterns FIRST to avoid substring ambiguity
+ # (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`).
+ stacked_params_mapping: ClassVar[list[tuple[str, str, str | int]]] = [
+ (".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"),
+ (".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"),
+ (".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"),
+ (".qkv_proj", ".q_proj", "q"),
+ (".qkv_proj", ".k_proj", "k"),
+ (".qkv_proj", ".v_proj", "v"),
+ # MLP gate/up fused into MergedColumnParallelLinear
+ (".gate_up_proj", ".gate_proj", 0),
+ (".gate_up_proj", ".up_proj", 1),
+ ]
+
+ def load_lora_weights(
+ self,
+ pretrained_model_name_or_path: str | list[str],
+ adapter_name: str | None = None,
+ ) -> None:
+ """Fuse a distilled few-step LoRA into the weights.
+
+ The checkpoints ship kohya `lora_down`/`lora_up`/`alpha` names, renamed
+ here to the Diffusers names before `load_lora_into_module` routes each
+ delta into its slice of the fused projections.
+ """
+ if adapter_name in self.lora_loaded:
+ return
+
+ paths = (
+ [pretrained_model_name_or_path]
+ if isinstance(pretrained_model_name_or_path, str)
+ else list(pretrained_model_name_or_path)
+ )
+ if len(paths) != 1:
+ raise ValueError(f"Expected exactly one distilled LoRA file, got {len(paths)}: {paths}")
+
+ state_dict = _load_lora_state_dict(paths[0])
+ state_dict = _remap_state_dict_keys(
+ state_dict,
+ [(".lora_down.weight", ".lora_A.weight"), (".lora_up.weight", ".lora_B.weight")],
+ )
+ # Fold alpha and the low-rank product in fp32; scaling B in bf16 would
+ # round once before the matmul and once after.
+ state_dict = {k: v.to(torch.float32) for k, v in state_dict.items()}
+ state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)
+
+ applied = self._add_lora_deltas(state_dict)
+ if not applied:
+ raise ValueError(f"Distilled LoRA {paths[0]} matched no parameter of this pipeline.")
+ logger.info("Fused distilled LoRA %s into %d parameters", paths[0], applied)
+ # Fusion is one-way and there is no unload path, so keep only a sentinel;
+ # the fp32 state dict is ~1.5 GiB and would stay resident for the process.
+ self.lora_loaded[adapter_name] = paths[0]
+
+ def _fused_shards(self, base_key: str) -> list[tuple[str, str | int]] | None:
+ """Unfused source names and shard ids for a fused parameter, else None."""
+ for param_name, weight_name, shard_id in self.stacked_params_mapping:
+ if base_key.endswith(param_name):
+ return [
+ (base_key[: -len(param_name)] + w, sid)
+ for p, w, sid in self.stacked_params_mapping
+ if p == param_name
+ ]
+ return None
+
+ def _add_lora_deltas(self, state_dict: dict[str, torch.Tensor]) -> int:
+ """Add every matching delta into its parameter, sharding through vLLM.
+
+ The delta read from the checkpoint spans the whole layer. A parameter of
+ a parallel layer only holds this rank's slice of it, so the delta goes
+ through the layer's own weight loader -- into a zeroed copy of the
+ parameter, which is then added -- instead of being added directly.
+
+ Each parameter is fused in place as it is reached, so a failure part way
+ through leaves the earlier ones fused and the pipeline has to be
+ reloaded. Deferring the writes would mean holding every delta at once,
+ and the deltas are full rank: 15.09 GiB for the 8B checkpoint, against
+ the 192 MiB one parameter costs here. Fusion runs at load time, before
+ the pipeline serves anything, so a partially fused model is never
+ reachable from a request.
+ """
+ applied = 0
+ for name, param in self.named_parameters():
+ if not name.endswith(".weight"):
+ continue
+ base_key = name[: -len(".weight")]
+ shards = self._fused_shards(base_key)
+
+ deltas: list[tuple[torch.Tensor, str | int | None]] = []
+ for src, shard_id in shards or [(base_key, None)]:
+ delta, _ = _prepare_lora_delta(state_dict, src, None)
+ if delta is None:
+ deltas = []
+ break
+ deltas.append((delta, shard_id))
+ if not deltas:
+ continue
+
+ weight_loader = getattr(param, "weight_loader", None)
+ buffer = torch.zeros_like(param.data)
+ if weight_loader is None:
+ # Not a parallel layer: the delta is already rank-local.
+ buffer = torch.cat([d for d, _ in deltas]).to(device=param.device, dtype=param.dtype)
+ else:
+ saved = param.data
+ try:
+ param.data = buffer
+ for delta, shard_id in deltas:
+ delta = delta.to(device=param.device, dtype=param.dtype)
+ if shard_id is None:
+ weight_loader(param, delta)
+ else:
+ weight_loader(param, delta, shard_id)
+ finally:
+ param.data = saved
+ with torch.no_grad():
+ param.data.add_(buffer)
+ applied += 1
+ return applied
+
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
- stacked_params_mapping = [
- # More specific _mot_gen patterns FIRST to avoid substring
- # ambiguity (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`).
- (".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"),
- (".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"),
- (".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"),
- (".qkv_proj", ".q_proj", "q"),
- (".qkv_proj", ".k_proj", "k"),
- (".qkv_proj", ".v_proj", "v"),
- # MLP gate/up fused into MergedColumnParallelLinear
- (".gate_up_proj", ".gate_proj", 0),
- (".gate_up_proj", ".up_proj", 1),
- ]
- self.stacked_params_mapping = stacked_params_mapping
+ stacked_params_mapping = self.stacked_params_mapping
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
index d2db1f489..ab69d3baa 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Qwen3 LLM with Mixture-of-Tokenizers (MoT) for SenseNova-U1.
Ported from the sensenova_u1 package with vllm tensor-parallel support:
@@ -16,6 +16,7 @@ from dataclasses import dataclass
import torch
import torch.nn as nn
+import torch.nn.functional as F
from cache_dit import ForwardPattern
from transformers.cache_utils import DynamicCache
from vllm.logger import init_logger
@@ -163,6 +164,18 @@ class Qwen3RotaryEmbedding(nn.Module):
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+def _build_3d_rope(config) -> tuple[Qwen3RotaryEmbedding, Qwen3RotaryEmbedding]:
+ """The t and h/w rotary modules: t uses half head_dim, h/w a quarter each."""
+ t_config = copy.deepcopy(config)
+ t_config.head_dim = config.head_dim // 2
+
+ hw_config = copy.deepcopy(config)
+ hw_config.head_dim = config.head_dim // 4
+ hw_config.rope_theta = config.rope_theta_hw
+ hw_config.max_position_embeddings = config.max_position_embeddings_hw
+ return Qwen3RotaryEmbedding(config=t_config), Qwen3RotaryEmbedding(config=hw_config)
+
+
# ---------------------------------------------------------------------------
# Norms
# ---------------------------------------------------------------------------
@@ -175,11 +188,9 @@ class Qwen3RMSNorm(nn.Module):
self.variance_epsilon = eps
def forward(self, hidden_states):
- input_dtype = hidden_states.dtype
- hidden_states = hidden_states.to(torch.float32)
- variance = hidden_states.pow(2).mean(-1, keepdim=True)
- hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
- return self.weight * hidden_states.to(input_dtype)
+ # F.rms_norm keeps the fp32 accumulation internal and rounds once; the
+ # cast chain it replaces rounded to bf16 before the weight multiply.
+ return F.rms_norm(hidden_states, self.weight.shape, self.weight, self.variance_epsilon)
# ---------------------------------------------------------------------------
@@ -303,23 +314,16 @@ class SenseNovaU1Attention(nn.Module):
self.q_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
self.k_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps)
- # 3D RoPE: t uses half head_dim, h/w each use quarter head_dim
- t_config = copy.deepcopy(config)
- t_config.head_dim = config.head_dim // 2
- self.rotary_emb = Qwen3RotaryEmbedding(config=t_config)
-
- hw_config = copy.deepcopy(config)
- hw_config.head_dim = config.head_dim // 4
- hw_config.rope_theta = config.rope_theta_hw
- hw_config.max_position_embeddings = config.max_position_embeddings_hw
- self.rotary_emb_hw = Qwen3RotaryEmbedding(config=hw_config)
-
self.attn = Attention(
num_heads=self.num_heads,
head_size=self.head_dim,
causal=False,
softmax_scale=self.scaling,
- num_kv_heads=self.num_heads,
+ # Keep K/V at their real head count. The SDPA backend already picks
+ # a fused GQA kernel when the runtime shape allows one and expands
+ # K/V itself when it does not; expanding here hides the GQA shape
+ # from that check, so it never fires.
+ num_kv_heads=self.num_kv_heads,
prefix=f"{prefix}.attn",
)
self.attn.attention = self.attn.sdpa_fallback
@@ -339,10 +343,6 @@ class SenseNovaU1Attention(nn.Module):
attention_mask: torch.Tensor | None,
) -> torch.Tensor:
"""Run unified attention with [B, H, S, D] inputs. Returns [B, S, H, D]."""
- if self.num_kv_groups > 1:
- n = self.num_kv_groups
- key_bhsd = key_bhsd.repeat_interleave(n, dim=1)
- value_bhsd = value_bhsd.repeat_interleave(n, dim=1)
q = query_bhsd.transpose(1, 2).contiguous()
k = key_bhsd.transpose(1, 2).contiguous()
v = value_bhsd.transpose(1, 2).contiguous()
@@ -358,16 +358,16 @@ class SenseNovaU1Attention(nn.Module):
attention_mask: torch.Tensor | None,
) -> torch.Tensor:
"""Run unified attention with [B, S, H, D] inputs. Returns [B, S, H, D]."""
- if self.num_kv_groups > 1:
- n = self.num_kv_groups
- key_bshd = key_bshd.repeat_interleave(n, dim=2)
- value_bshd = value_bshd.repeat_interleave(n, dim=2)
attention_mask = self._align_mask_dtype(attention_mask, query_bshd)
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
return self.attn(query_bshd, key_bshd, value_bshd, attn_metadata)
- def _project_and_rope(self, hidden_states, indexes, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw):
- """Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE."""
+ def _project_and_rope(self, hidden_states, position_embeddings, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw):
+ """Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE.
+
+ The three tables come from the model, which builds them once per forward
+ because every layer is handed the same `indexes`.
+ """
input_shape = hidden_states.shape[:-1] # (B, S)
qkv, _ = qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
@@ -376,9 +376,7 @@ class SenseNovaU1Attention(nn.Module):
k = k.view(*input_shape, self.num_kv_heads, self.head_dim)
v = v.view(*input_shape, self.num_kv_heads, self.head_dim)
- cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0))
- cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0))
- cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0))
+ (cos_t, sin_t), (cos_h, sin_h), (cos_w, sin_w) = position_embeddings
value_states = v.transpose(1, 2) # [B, H, S, D]
try:
@@ -433,13 +431,23 @@ class SenseNovaU1Attention(nn.Module):
input_shape = hidden_states.shape[:-1]
query_states, key_states, value_states = self._project_and_rope(
hidden_states,
- indexes,
+ kwargs["position_embeddings"],
self.qkv_proj,
self.q_norm,
self.k_norm,
self.q_norm_hw,
self.k_norm_hw,
)
+ # Single-token decode with a paged cache: append this step's K/V and let
+ # the kernel read the valid prefix via `seqused_k`. Shapes stay fixed,
+ # which is what makes the step capturable as a CUDA graph.
+ paged = kwargs.get("paged_cache")
+ if paged is not None and attention_mask is None and query_states.shape[2] == 1:
+ attn_output = paged.attend(self.layer_idx, query_states, key_states, value_states, self.scaling)
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output, _ = self.o_proj(attn_output)
+ return attn_output
+
update_cache = kwargs.get("update_cache", True)
if past_key_values is not None:
if update_cache:
@@ -460,7 +468,7 @@ class SenseNovaU1Attention(nn.Module):
input_shape = hidden_states.shape[:-1]
query_states, key_states, value_states = self._project_and_rope(
hidden_states,
- indexes,
+ kwargs["position_embeddings"],
self.qkv_proj_mot_gen,
self.q_norm_mot_gen,
self.k_norm_mot_gen,
@@ -659,6 +667,10 @@ class SenseNovaU1Model(nn.Module):
)
self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ # Every layer receives the same `indexes`, so the tables are the same for
+ # all of them: build them once per forward here and thread them down,
+ # rather than giving every layer its own pair.
+ self.rotary_emb, self.rotary_emb_hw = _build_3d_rope(config)
def forward(
self,
@@ -690,15 +702,25 @@ class SenseNovaU1Model(nn.Module):
past_len = past_key_values.get_seq_length() if past_key_values else 0
seq_len = inputs_embeds.shape[1]
total_len = past_len + seq_len
- mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device)
if seq_len > 1:
+ mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device)
causal = torch.tril(torch.ones(seq_len, seq_len, device=inputs_embeds.device))
mask[:, :, :, past_len:] = torch.where(causal == 1, 0.0, float("-inf"))
+ else:
+ # A single query token attends to every cached key, so the mask
+ # this branch used to build was all zeros. Passing it changed no
+ # result, only which SDPA kernel was reachable.
+ mask = None
causal_mask_mapping = {"full_attention": mask}
else:
causal_mask_mapping = attention_mask
hidden_states = inputs_embeds
+ position_embeddings = (
+ self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)),
+ self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)),
+ self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)),
+ )
for layer in self.layers:
hidden_states = layer(
hidden_states,
@@ -706,6 +728,7 @@ class SenseNovaU1Model(nn.Module):
exist_und=exist_und,
exist_gen=exist_gen,
indexes=indexes,
+ position_embeddings=position_embeddings,
attention_mask=causal_mask_mapping,
past_key_values=past_key_values,
use_cache=use_cache,
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index ea0dff90a..61f45c76e 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -519,6 +519,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
if self.prompt_embed_cache is not None:
self.prompt_embed_cache.clear()
+ def release_captured_graphs(self) -> None:
+ """Drop every CUDA graph held for this model, wherever it is kept.
+
+ Sleep level 2 discards the memory a capture was recorded against, so a
+ graph that outlives it replays over freed storage. The runner owns
+ compilation and execution resources, so it owns the release: a pipeline
+ that captures graphs of its own implements ``release_captured_graphs``
+ and is collected here, instead of every caller having to know which
+ pipelines have one.
+ """
+ runners = getattr(self, "graph_runners", None)
+ if runners is not None:
+ runners.clear()
+ release = getattr(self.pipeline, "release_captured_graphs", None)
+ if callable(release):
+ release()
+
def get_prompt_embed_cache_stats(self) -> dict | None:
"""Return hit/miss statistics for the prompt-embedding cache, if enabled.
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 5f8fa3e82..033848008 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -786,9 +786,8 @@ class DiffusionWorker:
usage_before = allocator.get_current_usage()
if level == 2 and self.model_runner is not None:
- if hasattr(self.model_runner, "graph_runners"):
- self.model_runner.graph_runners.clear()
- logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
+ self.model_runner.release_captured_graphs()
+ logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
model = self.model_runner.pipeline
self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}