[Attention] Add trtllm diffusion attention backend with Skip-Softmax (#5283)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

tests/diffusion/attention/test_attention_config.py

@@ -31,21 +31,27 @@ from vllm_omni.diffusion.data import (
     parse_attention_config,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 class TestAttentionSpec:
-    def test_construct_no_extra(self):
+    def test_construct_no_skip_softmax(self):
         spec = AttentionSpec(backend="FLASH_ATTN")
-        assert spec.extra == {}
+        assert spec.skip_softmax is None
 
-    def test_mapping_extra_normalized(self):
-        spec = AttentionSpec(backend="SAGE_ATTN", extra={"quant": "int8"})
-        assert spec.backend == "SAGE_ATTN"
-        assert spec.extra == {"quant": "int8"}
+    def test_skip_softmax_mapping_coerced(self):
+        spec = AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})
+        assert spec.backend == "TRTLLM_ATTN"
+        assert spec.skip_softmax.target_sparsity == 0.5
 
     def test_invalid_backend_type(self):
         with pytest.raises(TypeError):
             AttentionSpec(backend=123)  # type: ignore[arg-type]
 
+    def test_skip_softmax_rejected_on_non_trtllm(self):
+        with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"):
+            AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5})
+
 
 class TestAttentionConfig:
     def test_empty_config(self):
@@ -57,13 +63,13 @@ class TestAttentionConfig:
         config = AttentionConfig(
             default={"backend": "FLASH_ATTN"},
             per_role={
-                "self": {"backend": "SPARSE_BLOCK", "extra": {"block_size": 128}},
+                "self": {"backend": "TRTLLM_ATTN", "skip_softmax": {"target_sparsity": 0.5}},
                 "cross": "SAGE_ATTN",
             },
         )
         assert config.default.backend == "FLASH_ATTN"
-        assert config.per_role["self"].backend == "SPARSE_BLOCK"
-        assert config.per_role["self"].extra == {"block_size": 128}
+        assert config.per_role["self"].backend == "TRTLLM_ATTN"
+        assert config.per_role["self"].skip_softmax.target_sparsity == 0.5
         assert config.per_role["cross"].backend == "SAGE_ATTN"
 
     def test_constructor_flattens_nested_per_role_tree(self):
@@ -178,10 +184,10 @@ class TestAttentionConfig:
         config = AttentionConfig(
             default=AttentionSpec(backend="FLASH_ATTN"),
             per_role={
-                "self": AttentionSpec(backend="SPARSE_BLOCK", extra={"block_size": 128}),
+                "self": AttentionSpec(backend="SPARSE_BLOCK"),
                 "cross": AttentionSpec(backend="SAGE_ATTN"),
                 "ltx2.audio_self": AttentionSpec(backend="FLASH_ATTN"),
-                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN", extra={"causal_window": 64}),
+                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN"),
             },
         )
 
@@ -200,7 +206,6 @@ class TestAttentionConfig:
         # audio-to-video → exact match
         spec, _ = config.resolve_with_source("ltx2.audio_to_video", "cross")
         assert spec.backend == "FLASH_ATTN"
-        assert spec.extra == {"causal_window": 64}
 
         # video-to-audio → category fallback to "cross"
         assert config.resolve_with_source("ltx2.video_to_audio", "cross")[0].backend == "SAGE_ATTN"
@@ -364,12 +369,13 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
             head_size,
             attention_config=None,
             role_category=None,
+            allow_trtllm_default=False,
         ):
             captured["role"] = role
             captured["head_size"] = head_size
             captured["role_category"] = role_category
             captured["attention_config"] = attention_config
-            return _FakeBackend, AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})
+            return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})
 
         class _FakeRingParallelAttention:
             def __init__(self, sp_group, attn_backend_pref=None):
@@ -391,7 +397,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         od_config = SimpleNamespace(
             diffusion_attention_config=AttentionConfig(
                 default=AttentionSpec(backend="FLASH_ATTN"),
-                per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
+                per_role={"cross": AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})},
             ),
             parallel_config=SimpleNamespace(ring_degree=2),
             diffusion_kv_cache_dtype=None,
@@ -414,12 +420,12 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         assert captured["role_category"] == "cross"
         assert captured["head_size"] == 64
         assert captured["attention_config"] is od_config.diffusion_attention_config
-        assert attn.backend_pref == "TORCH_SDPA"
-        assert attn.attention.kwargs["backend_kwargs"] == {"block_size": 128}
+        assert attn.backend_pref == "TRTLLM_ATTN"
+        assert attn.attention.kwargs["backend_kwargs"] == {"target_sparsity": 0.5}
         assert attn.attention.kwargs["qkv_layout"] == "BSND"
         assert attn.use_ring is True
         assert attn.ring_runner is not None
-        assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"
+        assert attn.ring_runner.attn_backend_pref == "TRTLLM_ATTN"
 
 
 class TestDiffusionKvCacheQuantization:
@@ -453,7 +459,10 @@ class TestDiffusionKvCacheQuantization:
         monkeypatch.setattr(
             layer_mod,
             "get_attn_backend_for_role",
-            lambda role, head_size, attention_config=None, role_category=None: (_FakeBackend, None),
+            lambda role, head_size, attention_config=None, role_category=None, allow_trtllm_default=False: (
+                _FakeBackend,
+                None,
+            ),
         )
         monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
         monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())

tests/diffusion/attention/test_trtllm_attn.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
+from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+def _has_trtllm_attn() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    if torch.cuda.get_device_capability()[0] < 10:
+        return False
+    try:
+        import flashinfer  # noqa: F401
+    except Exception:
+        return False
+    return tg.HAS_FLASHINFER
+
+
+requires_trtllm_attn = pytest.mark.skipif(
+    not _has_trtllm_attn(), reason="requires Blackwell SM100+ GPU with flashinfer"
+)
+
+
+def _impl(**bk):
+    return TrtllmAttentionImpl(8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=bk)
+
+
+def test_skip_config_pure_resolution():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
+
+    assert not SkipSoftmaxConfig().enabled
+    assert SkipSoftmaxConfig().resolve_factor(4096, timestep=0.3) is None
+
+    cfg = SkipSoftmaxConfig.from_backend_kwargs({"skip_softmax_threshold": 0.01, "disabled_until_timestep": 0.6})
+    assert cfg.enabled and cfg.gated
+    assert cfg.resolve_factor(4096, timestep=0.9) is None
+    assert cfg.resolve_factor(4096, timestep=0.3) == pytest.approx(0.01 * 4096)
+    assert cfg.resolve_factor(4096, timestep=None) == pytest.approx(0.01 * 4096)
+
+
+def test_skip_factor_none_without_curve():
+    assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
+
+
+def _calibrated(a, b, **bk):
+    impl = _impl(**bk)
+    impl.set_layer_calibration(a, b)
+    return impl
+
+
+def test_skip_factor_from_curve():
+    assert _calibrated(100.0, 0.0, target_sparsity=0.5)._resolve_skip_factor(4096) == pytest.approx(100.0)
+    f2 = _calibrated(2.0, 1.0, target_sparsity=0.5)._resolve_skip_factor(4096)
+    assert f2 == pytest.approx(2.0 * math.exp(0.5))
+
+
+def test_skip_factor_none_until_stamped():
+    assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
+
+
+def test_skip_factor_from_threshold_no_calibration():
+    impl = _impl(skip_softmax_threshold=0.01)
+    assert impl._resolve_skip_factor(4096) == pytest.approx(0.01 * 4096)
+    assert impl._resolve_skip_factor(8192) == pytest.approx(0.01 * 8192)
+
+    impl2 = _calibrated(100.0, 0.0, skip_softmax_threshold=0.02, target_sparsity=0.5)
+    assert impl2._resolve_skip_factor(4096) == pytest.approx(0.02 * 4096)
+
+
+def test_skip_timestep_gate(monkeypatch):
+    from vllm_omni.diffusion import forward_context as fc
+
+    impl = _calibrated(100.0, 0.0, target_sparsity=0.5, disabled_until_timestep=0.6)
+    ctx = fc.ForwardContext(denoise_timestep=0.9)
+    monkeypatch.setattr(fc, "_forward_context", ctx)
+    assert impl._resolve_skip_factor(4096) is None
+    ctx.denoise_timestep = 0.3
+    assert impl._resolve_skip_factor(4096) == pytest.approx(100.0)
+
+
+def test_denoise_progress_mixin(monkeypatch):
+    import types
+
+    from vllm_omni.diffusion import forward_context as fc
+
+    class _Pipe(fc.DenoiseProgressMixin):
+        scheduler = types.SimpleNamespace(config=types.SimpleNamespace(num_train_timesteps=1000))
+
+    ctx = fc.ForwardContext()
+    monkeypatch.setattr(fc, "_forward_context", ctx)
+    _Pipe().record_denoise_step(3, 250)
+    assert ctx.denoise_step_idx == 3
+    assert ctx.denoise_timestep == pytest.approx(0.25)
+
+    _Pipe().record_denoise_step(4)
+    assert ctx.denoise_step_idx == 4 and ctx.denoise_timestep == pytest.approx(0.25)
+
+
+@pytest.mark.parametrize(
+    "bk, field",
+    [
+        ({"target_sparsity": 1.5}, "target_sparsity"),
+        ({"target_sparsity": -0.1}, "target_sparsity"),
+        ({"disabled_until_timestep": 1.2}, "disabled_until_timestep"),
+        ({"skip_softmax_threshold": -1.0}, "skip_softmax_threshold"),
+        ({"target_sparsity": float("nan")}, "target_sparsity"),
+        ({"skip_softmax_threshold": float("inf")}, "skip_softmax_threshold"),
+    ],
+)
+def test_skip_control_validation_rejects_out_of_range(bk, field):
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
+
+    with pytest.raises(ValueError, match=field):
+        SkipSoftmaxConfig.from_backend_kwargs(bk)
+
+
+def test_skip_control_validation_accepts_valid():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
+
+    cfg = SkipSoftmaxConfig.from_backend_kwargs(
+        {"target_sparsity": 0.65, "disabled_until_timestep": 0.86, "skip_softmax_threshold": 0.01}
+    )
+    assert cfg.target_sparsity == pytest.approx(0.65)
+    assert cfg.disabled_until_timestep == pytest.approx(0.86)
+    assert cfg.threshold == pytest.approx(0.01)
+
+
+def test_ring_sp_with_skip_softmax_rejected():
+    import types
+
+    from vllm_omni.diffusion.attention.layer import Attention
+
+    fake = types.SimpleNamespace(attention=_impl(target_sparsity=0.5), ring_runner=None)
+    with pytest.raises(NotImplementedError, match="ring sequence parallelism"):
+        Attention._run_ring_attention(fake, None, None, None, None)
+
+
+def test_mask_free_allowlist_gates_trtllm_default():
+    from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata as meta
+
+    assert meta("WanPipeline").attention_mask_free is True
+    assert meta("WanVACEPipeline").attention_mask_free is True
+    assert meta("FluxPipeline").attention_mask_free is False
+    assert meta(None).attention_mask_free is False
+
+
+def test_masked_layer_raises():
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    mask = torch.ones(b, s, dtype=torch.bool)
+    mask[:, s // 2 :] = False
+
+    class _Meta:
+        attn_mask = mask
+
+    with pytest.raises(ValueError, match="does not support attn_mask"):
+        _impl().forward_cuda(q, k, v, _Meta())
+
+
+def test_propagate_calibration_045_schema(monkeypatch):
+    import types
+
+    from vllm_omni.diffusion.data import (
+        AttentionConfig,
+        AttentionSpec,
+        OmniDiffusionConfig,
+        TransformerConfig,
+    )
+
+    cfg = AttentionConfig(default=AttentionSpec(backend="TRTLLM_ATTN"))
+    tf = TransformerConfig.from_dict(
+        {
+            "sparse_attention_config": {
+                "config_groups": {
+                    "group_0": {
+                        "algorithm": "skip_softmax",
+                        "targets": ["WanAttention"],
+                        "ignore": ["blocks.0.attn2", "blocks.1.attn2"],
+                        "threshold_scale_factor": {
+                            "formula": "a * exp(b * target_sparsity)",
+                            "coefficients": {"a": 2142.7, "b": 4.28},
+                        },
+                        "target_sparsity": 0.5,
+                    }
+                }
+            }
+        }
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.trtllm_calibration.parse_secondary_expert_calibration",
+        lambda *a, **k: None,
+    )
+    OmniDiffusionConfig._propagate_skip_softmax_calibration(
+        types.SimpleNamespace(
+            diffusion_attention_config=cfg,
+            model="/x",
+        ),
+        tf,
+    )
+    spec, _ = cfg.resolve_with_source("self")
+    cal = spec.skip_calibration["by_expert"]["transformer"]
+    assert cal["a"] == pytest.approx(2142.7)
+    assert cal["b"] == pytest.approx(4.28)
+    assert cal["ignore"] == ["blocks.0.attn2", "blocks.1.attn2"]
+    assert "transformer_2" not in spec.skip_calibration["by_expert"]
+
+
+def test_target_sparsity_without_calibration_raises():
+    import types
+
+    from vllm_omni.diffusion.data import (
+        AttentionConfig,
+        AttentionSpec,
+        OmniDiffusionConfig,
+        TransformerConfig,
+    )
+
+    def _cfg(skip_softmax):
+        return types.SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(
+                default=AttentionSpec(backend="TRTLLM_ATTN", skip_softmax=skip_softmax)
+            ),
+            model="/models/no-such-calibration",
+        )
+
+    with pytest.raises(ValueError, match="no skip-softmax calibration"):
+        OmniDiffusionConfig._propagate_skip_softmax_calibration(
+            _cfg({"target_sparsity": 0.5}), TransformerConfig.from_dict({})
+        )
+
+    OmniDiffusionConfig._propagate_skip_softmax_calibration(_cfg({"threshold": 0.02}), TransformerConfig.from_dict({}))
+
+
+def _sdpa_ref(q, k, v, scale):
+    qp, kp, vp = (x.permute(0, 2, 1, 3) for x in (q, k, v))
+    return F.scaled_dot_product_attention(qp, kp, vp, is_causal=False, scale=scale).permute(0, 2, 1, 3).float()
+
+
+@requires_trtllm_attn
+def test_bf16_dense_matches_sdpa():
+    torch.manual_seed(0)
+    B, S, H, D = 2, 256, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    out = _impl().forward_cuda(q, k, v, None).float()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
+
+
+@requires_trtllm_attn
+def test_skip_tiny_threshold_matches_dense():
+    torch.manual_seed(0)
+    B, S, H, D = 2, 512, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+
+    out = _impl(skip_softmax_threshold=1e-9).forward_cuda(q, k, v, None).float()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.02, f"skip(tiny) rel err {rel:.4f} should be ~dense"
+
+
+@requires_trtllm_attn
+def test_skip_threshold_engages():
+    torch.manual_seed(0)
+    B, S, H, D = 2, 512, 8, 128
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    dense = _impl().forward_cuda(q, k, v, None).float()
+    skipped = _impl(skip_softmax_threshold=2.0).forward_cuda(q, k, v, None).float()
+    assert torch.isfinite(skipped).all()
+    assert (skipped - dense).abs().mean() / dense.abs().mean() > 0.02
+
+
+@requires_trtllm_attn
+def test_gqa_matches_sdpa():
+    torch.manual_seed(0)
+    B, S, Hq, Hkv, D = 1, 256, 64, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, S, Hq, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(B, S, Hkv, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(B, S, Hkv, D, device="cuda", dtype=torch.bfloat16)
+    out = (
+        TrtllmAttentionImpl(Hq, D, scale, causal=False, num_kv_heads=Hkv, backend_kwargs={})
+        .forward_cuda(q, k, v, None)
+        .float()
+    )
+    qp, kp, vp = (x.permute(0, 2, 1, 3) for x in (q, k, v))
+    ref = (
+        F.scaled_dot_product_attention(qp, kp, vp, is_causal=False, scale=scale, enable_gqa=True)
+        .permute(0, 2, 1, 3)
+        .float()
+    )
+    rel = (out - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.01, f"GQA rel err {rel:.4f} vs SDPA(GQA) too high"
+
+
+@requires_trtllm_attn
+def test_unsupported_head_dim_raises():
+    q = k = v = torch.randn(1, 32, 8, 64, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(Exception, match="[Uu]nsupported head dim"):
+        _impl().forward_cuda(q, k, v, None)
+
+
+def test_skip_end_to_end_config_path(monkeypatch):
+    import types
+
+    from vllm_omni.diffusion import forward_context as fc
+    from vllm_omni.diffusion.data import (
+        AttentionConfig,
+        AttentionSpec,
+        OmniDiffusionConfig,
+        TransformerConfig,
+    )
+
+    cfg = AttentionConfig(
+        default=AttentionSpec(
+            backend="TRTLLM_ATTN",
+            skip_softmax={"target_sparsity": 0.65, "disabled_until_timestep": 0.86},
+        )
+    )
+
+    tf = TransformerConfig.from_dict(
+        {"sparse_attention_config": {"threshold_scale_factor": {"prefill": {"a": 7.93, "b": 8.61}}}}
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.trtllm_calibration.parse_secondary_expert_calibration",
+        lambda *a, **k: None,
+    )
+    OmniDiffusionConfig._propagate_skip_softmax_calibration(
+        types.SimpleNamespace(
+            diffusion_attention_config=cfg,
+            model="/x",
+        ),
+        tf,
+    )
+    spec, _ = cfg.resolve_with_source("self")
+    cal = spec.skip_calibration["by_expert"]["transformer"]
+    assert cal["a"] == 7.93 and cal["b"] == 8.61
+
+    impl = TrtllmAttentionImpl(
+        8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=spec.backend_kwargs()
+    )
+    impl.set_layer_calibration(cal["a"], cal["b"])
+    expected = 7.93 * math.exp(8.61 * 0.65)
+
+    ctx = fc.ForwardContext(denoise_timestep=0.3)
+    monkeypatch.setattr(fc, "_forward_context", ctx)
+    assert impl._resolve_skip_factor(4096) == pytest.approx(expected)

tests/diffusion/attention/test_trtllm_calibration.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
+    is_ignored,
+    layer_match_names,
+    parse_sparse_attention_config,
+    resolve_layer_calibration,
+    select_expert,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_layer_match_names_strips_expert_prefix():
+    names = layer_match_names("transformer_2.blocks.5.attn1")
+    assert "transformer_2.blocks.5.attn1" in names
+    assert "blocks.5.attn1" in names
+
+
+def test_layer_match_names_strips_orig_mod():
+    names = layer_match_names("transformer.blocks.5._orig_mod.attn1")
+    assert "transformer.blocks.5.attn1" in names or "blocks.5.attn1" in names
+
+
+def test_is_ignored_matches_relative_pattern():
+    assert is_ignored("transformer.blocks.5.attn2", ["blocks.*.attn2"])
+    assert not is_ignored("transformer.blocks.5.attn1", ["blocks.*.attn2"])
+
+
+def test_is_ignored_empty_patterns():
+    assert not is_ignored("transformer.blocks.5.attn1", [])
+    assert not is_ignored("transformer.blocks.5.attn1", None)
+
+
+def test_is_ignored_ancestor_match_with_attn_suffix():
+    assert is_ignored("transformer.blocks.0.attn2.attn", ["blocks.0.attn2"])
+
+    assert not is_ignored("transformer.blocks.11.attn2.attn", ["blocks.1.attn2"])
+
+    assert not is_ignored("transformer.blocks.5.attn1.attn", ["blocks.0.attn2", "blocks.1.attn2"])
+
+
+def test_select_expert_longest_prefix_wins():
+    keys = ("transformer", "transformer_2")
+    assert select_expert("transformer_2.blocks.5.attn1", keys) == "transformer_2"
+    assert select_expert("transformer.blocks.5.attn1", keys) == "transformer"
+
+
+def test_select_expert_no_match():
+    assert select_expert("blocks.5.attn1", ("transformer", "transformer_2")) is None
+
+
+def _two_expert_calib():
+    return {
+        "by_expert": {
+            "transformer": {
+                "a": 2142.7,
+                "b": 4.28,
+                "target_sparsity": 0.45,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2"],
+            },
+            "transformer_2": {
+                "a": 314.68,
+                "b": 6.17,
+                "target_sparsity": 0.45,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2"],
+            },
+        }
+    }
+
+
+def test_resolve_routes_to_correct_expert():
+    calib = _two_expert_calib()
+    hi = resolve_layer_calibration("transformer.blocks.5.attn1", calib)
+    lo = resolve_layer_calibration("transformer_2.blocks.5.attn1", calib)
+    assert hi["a"] == pytest.approx(2142.7)
+    assert lo["a"] == pytest.approx(314.68)
+
+
+def test_resolve_ignored_layer_is_dense():
+    assert resolve_layer_calibration("transformer.blocks.5.attn2", _two_expert_calib()) is None
+
+
+def test_resolve_single_transformer():
+    calib = {
+        "a": 104.76,
+        "b": 7.81,
+        "target_sparsity": 0.45,
+        "formula": "a*exp(b*target_sparsity)",
+        "ignore": ["refiner_blocks.*"],
+    }
+    assert resolve_layer_calibration("transformer_blocks.5.attn", calib)["a"] == pytest.approx(104.76)
+    assert resolve_layer_calibration("refiner_blocks.0.attn", calib) is None
+
+
+def test_resolve_none_calibration():
+    assert resolve_layer_calibration("transformer.blocks.5.attn1", None) is None
+
+
+def test_resolve_unmatched_component_is_uncalibrated():
+    calib = {"by_expert": {"transformer": {"a": 1.0, "b": 2.0, "ignore": []}}}
+    assert resolve_layer_calibration("transformer.blocks.5.attn1", calib)["a"] == pytest.approx(1.0)
+    assert resolve_layer_calibration("transformer_2.blocks.5.attn1", calib) is None
+    assert resolve_layer_calibration("blocks.5.attn1", calib) is None
+
+
+def test_parse_045_config_groups_schema():
+    sac = {
+        "config_groups": {
+            "group_0": {
+                "algorithm": "skip_softmax",
+                "ignore": ["blocks.0.attn1"],
+                "target_sparsity": 0.45,
+                "threshold_scale_factor": {
+                    "formula": "a * exp(b * target_sparsity)",
+                    "coefficients": {"a": 2142.7, "b": 4.28},
+                },
+            }
+        }
+    }
+    parsed = parse_sparse_attention_config(sac)
+    assert parsed["a"] == pytest.approx(2142.7)
+    assert parsed["b"] == pytest.approx(4.28)
+    assert parsed["ignore"] == ["blocks.0.attn1"]
+
+
+def test_parse_unsupported_formula_raises():
+    sac = {
+        "config_groups": {
+            "group_0": {
+                "algorithm": "skip_softmax",
+                "threshold_scale_factor": {"formula": "a * b ** target_sparsity", "coefficients": {"a": 1.0, "b": 2.0}},
+            }
+        }
+    }
+    with pytest.raises(ValueError, match="unsupported skip-softmax formula"):
+        parse_sparse_attention_config(sac)
+
+
+def test_parse_empty_returns_none():
+    assert parse_sparse_attention_config(None) is None
+    assert parse_sparse_attention_config({}) is None
+
+
+class _FakeImpl:
+    def __init__(self):
+        self.stamped = None
+
+    def set_layer_calibration(self, a, b):
+        self.stamped = (a, b)
+
+
+def _fake_pipeline():
+    import torch.nn as nn
+
+    def _attn_layer():
+        layer = nn.Module()
+        layer.attention = _FakeImpl()
+        return layer
+
+    def _wan_attn():
+        wa = nn.Module()
+        wa.attn = _attn_layer()
+        return wa
+
+    def _transformer():
+        tf = nn.Module()
+        blocks = nn.ModuleList()
+        for _ in range(3):
+            blk = nn.Module()
+            blk.attn1 = _wan_attn()
+            blk.attn2 = _wan_attn()
+            blocks.append(blk)
+        tf.blocks = blocks
+        return tf
+
+    pipe = nn.Module()
+    pipe.transformer = _transformer()
+    pipe.transformer_2 = _transformer()
+    return pipe
+
+
+def test_apply_to_pipeline_routes_and_ignores():
+    from vllm_omni.diffusion.attention.backends.trtllm_calibration import apply_to_pipeline
+
+    pipe = _fake_pipeline()
+
+    calib = {
+        "by_expert": {
+            "transformer": {
+                "a": 2142.7,
+                "b": 4.28,
+                "target_sparsity": 0.5,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2", "blocks.0.attn1"],
+            },
+            "transformer_2": {
+                "a": 314.68,
+                "b": 6.17,
+                "target_sparsity": 0.5,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2", "blocks.0.attn1"],
+            },
+        }
+    }
+    stamped = apply_to_pipeline(pipe, calib)
+
+    assert stamped == 4
+
+    def impl(expert, blk, attn):
+        return getattr(getattr(pipe, expert).blocks[blk], attn).attn.attention
+
+    assert impl("transformer", 1, "attn1").stamped == pytest.approx((2142.7, 4.28))
+
+    assert impl("transformer_2", 1, "attn1").stamped == pytest.approx((314.68, 6.17))
+
+    assert impl("transformer", 1, "attn2").stamped is None
+
+    assert impl("transformer", 0, "attn1").stamped is None