Update vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/host_weight_runtime.md

@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
 
 ## Status
 
-The first implementation provides contracts and a CPU local-filesystem store.
+The implementation provides contracts and a CPU local-filesystem store. The
-The initial diffusion consumer contract additionally defines typed,
+diffusion consumer additionally defines typed,
 representation-independent final-layout identity/restoration mechanics plus a
-concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains
+concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
-library-only: no loader or DLO path selects, publishes, restores, or transports
+`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
-that artifact yet.
+selects, publishes, restores, and transfers these artifacts.
 
 V1 includes:
 
@@ -33,13 +33,12 @@ V1 includes:
 
 V1 does not include:
 
-- a public CLI, loader activation, or default-on model integration;
+- default-on activation or consumers outside no-AllGather DLO;
-- lease handoff to DLO or another transport consumer;
+- online FP8, quantized, merged-adaptation, or additional model producers;
-- FP8, quantized, merged-adaptation, or additional model producers;
+- HWR interaction with DLO AllGather;
-- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
 - a remote artifact provider or cross-node coordination;
 - automatic eviction; or
-- a change to DLO AllGather or no-AllGather behavior.
+- a change to DLO collective or execution behavior.
 
 ## Motivation and use cases
 
@@ -255,7 +254,8 @@ The contract is intentionally separate from loader activation:
 - `FinalLayoutBF16Producer` accepts only the matching identity context and a
   finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
   coordinate. Its BF16 policy preserves model-declared FP32 parameters and
-  buffers and revalidates MiniMax H3 mixed-precision invariants.
+  buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
+  FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
 
 Other representations reuse source identity, typed parallel identity, tensor
 ownership, and exact restoration only when their policy proves those semantics.

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -220,9 +220,15 @@ model revision and parallel layout with `required`. TP coordinates have distinct
 identities, while equivalent DP replicas share them.
 
 The current producer/consumer boundary is the model-declared final-layout BF16
-contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1
+contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
-and TP2 rank identities plus SP layout identities; online quantization, HSDP,
+contract covers both transformer block stacks, constructor-stable packed QKV
-LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
+mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
+buffers. The shared HWR machinery supports ordinary-loader final layouts for
+TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
+in this change is TP1-only; its TP2/SP layouts rely on the existing
+per-coordinate identity mechanics and remain unmeasured. Online quantization,
+HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
+HWR mode
 `disabled`, DLO-disabled, and DLO AllGather configurations stop before source
 identity or store construction and retain the existing checkpoint-mmap or
 ordinary-loader path.

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -112,6 +112,12 @@ perform a collective and uses the same rank-local transfer behavior.
 HWR is an opt-in startup optimization for models that declare the final-layout
 BF16 restore contract. Use it only with no-AllGather DLO:
 
+The validated BF16 model contracts currently cover MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
+class but has not been validated against this contract. FLUX.2-dev, online
+FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
+the validated scope.
+
 ```bash
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -249,6 +249,30 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
     assert loader.take_host_weight_plan() is None
 
 
+def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
+    tmp_path: Path,
+    monkeypatch,
+):
+    canonical_root = tmp_path / "canonical"
+    canonical_root.mkdir()
+    save_file(
+        {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
+        str(canonical_root / "model.safetensors"),
+    )
+    loader = DiffusersPipelineLoader(
+        LoadConfig(),
+        _hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
+    )
+    pipeline = _HWRPipeline(canonical_root)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+
+    with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
+        loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    assert pipeline.load_count == 0
+    assert loader.take_host_weight_plan() is None
+
+
 def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     return OmniDiffusionConfig(
         model="",
@@ -261,10 +285,11 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
 
 
 @pytest.mark.parametrize(
-    ("dist_offload", "use_allgather"),
+    ("dist_offload", "use_allgather", "mode"),
     [
-        (False, False),
+        (False, False, "preferred"),
-        (True, True),
+        (True, True, "preferred"),
+        (True, False, "disabled"),
     ],
 )
 def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
@@ -272,12 +297,13 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
     tmp_path,
     dist_offload,
     use_allgather,
+    mode,
 ):
     """Disabled and AllGather paths must never construct or probe HWR."""
     from vllm_omni.host_weight_runtime import HostWeightRuntime
 
     root = tmp_path / "must-not-be-touched"
-    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root))
+    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
     model = _DummyPipelineModel(source_prefix="transformer.")
     modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
 

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -303,6 +303,83 @@ class _FakeH3Attention(nn.Module):
         super().__init__()
 
 
+class _FakeFluxLinear(_FakeH3Linear):
+    def __init__(
+        self,
+        *args: object,
+        bias: bool = False,
+        params_dtype: torch.dtype | None = None,
+        total_num_heads: int | None = None,
+        total_num_kv_heads: int | None = None,
+        **kwargs: object,
+    ) -> None:
+        super().__init__(
+            *args,
+            bias=bias,
+            params_dtype=params_dtype,
+            total_num_heads=total_num_heads,
+            total_num_kv_heads=total_num_kv_heads,
+            **kwargs,
+        )
+        self.weight.weight_loader = self._load_weight
+
+    @staticmethod
+    def _load_weight(
+        parameter: nn.Parameter,
+        loaded_weight: torch.Tensor,
+        shard_id: str | None = None,
+    ) -> None:
+        del shard_id
+        parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
+
+
+class _FakeFluxAttention(nn.Module):
+    def __init__(self, **kwargs: object) -> None:
+        del kwargs
+        super().__init__()
+        self.register_buffer("beta", torch.tensor([0.5]))
+        self.register_buffer("eps", torch.tensor([-1e-6]))
+
+
+class _Flux2KleinPipeline(nn.Module):
+    def __init__(self, transformer: nn.Module) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
+        self.vae = nn.Module()
+        self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
+
+
+def _reduced_flux2_klein(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    device: torch.device | str = "cpu",
+) -> nn.Module:
+    from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
+
+    monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
+    monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
+
+    transformer = flux2.Flux2Transformer2DModel(
+        patch_size=1,
+        in_channels=8,
+        out_channels=8,
+        num_layers=1,
+        num_single_layers=1,
+        attention_head_dim=8,
+        num_attention_heads=1,
+        joint_attention_dim=8,
+        timestep_guidance_channels=8,
+        mlp_ratio=1.0,
+        axes_dims_rope=(2, 2, 2, 2),
+        guidance_embeds=False,
+    )
+    return transformer.to(device=device, dtype=torch.bfloat16)
+
+
 class _H3Pipeline(nn.Module):
     def __init__(self, transformer: nn.Module) -> None:
         super().__init__()
@@ -1200,3 +1277,152 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
     assert contract.implementation_id == "minimax-h3-dit"
     assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
     assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
+
+
+def test_reduced_flux2_klein_declares_complete_final_layout_contract(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
+        Flux2Transformer2DModel,
+    )
+
+    transformer = _reduced_flux2_klein(monkeypatch)
+    pipeline = _Flux2KleinPipeline(transformer)
+    source = _prepared_source(tmp_path, directory="flux2-klein-source")
+    context = _identity(pipeline, source)
+    targets = collect_final_layout_targets(
+        pipeline,
+        (("transformer", transformer),),
+        policy=FINAL_LAYOUT_BF16_POLICY,
+        require_materialized=True,
+    )
+    target_names = {target.name for target in targets}
+
+    contract = Flux2Transformer2DModel.host_weight_restore_contract
+    assert isinstance(contract, FinalLayoutModelContract)
+    assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+    assert contract.implementation_id == "flux2-klein-dit"
+    assert contract.version == "1"
+    assert context.model_contracts == (contract,)
+    assert transformer.stacked_params_mapping == [
+        (".to_qkv.", ".to_q.", "q"),
+        (".to_qkv.", ".to_k.", "k"),
+        (".to_qkv.", ".to_v.", "v"),
+        (".add_kv_proj", ".add_q_proj", "q"),
+        (".add_kv_proj", ".add_k_proj", "k"),
+        (".add_kv_proj", ".add_v_proj", "v"),
+    ]
+    loaded = transformer.load_weights(
+        [
+            ("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
+            ("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
+            (
+                "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+                torch.tensor([3.0], dtype=torch.bfloat16),
+            ),
+        ]
+    )
+    assert loaded == {
+        "transformer_blocks.0.attn.to_qkv.weight",
+        "transformer_blocks.0.attn.add_kv_proj.weight",
+        "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+    }
+    assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
+    assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
+    assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
+    assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
+    assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
+    assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
+    assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
+    assert all("text_encoder" not in name and "vae" not in name for name in target_names)
+    transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_restore_validator_rejects_non_ready_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    transformer = _reduced_flux2_klein(monkeypatch)
+
+    mapping = transformer.stacked_params_mapping
+    transformer.stacked_params_mapping = None
+    with pytest.raises(ValueError, match="packed QKV mapping"):
+        transformer.validate_restored_host_weights()
+    transformer.stacked_params_mapping = mapping
+
+    weight = transformer.x_embedder.weight
+    weight.data = weight.data.float()
+    with pytest.raises(ValueError, match="must stay bf16"):
+        transformer.validate_restored_host_weights()
+    weight.data = weight.data.to(torch.bfloat16)
+
+    weight.data = weight.data.T
+    with pytest.raises(ValueError, match="contiguous strided layout"):
+        transformer.validate_restored_host_weights()
+    weight.data = weight.data.contiguous()
+
+    attention = transformer.transformer_blocks[0].attn.attn
+    attention._non_persistent_buffers_set.add("beta")
+    with pytest.raises(ValueError, match="must be persistent"):
+        transformer.validate_restored_host_weights()
+    attention._non_persistent_buffers_set.remove("beta")
+
+    attention.beta.data = attention.beta.data.float()
+    with pytest.raises(ValueError, match="scalar bf16"):
+        transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_publication_restore_fallback_and_teardown(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    source = _prepared_source(tmp_path, directory="flux2-klein-source")
+    cold_transformer = _reduced_flux2_klein(monkeypatch)
+    cold_model = _Flux2KleinPipeline(cold_transformer)
+    with torch.no_grad():
+        for index, parameter in enumerate(cold_transformer.parameters(), start=1):
+            parameter.fill_(index)
+        for name, buffer in cold_transformer.named_buffers():
+            if name.endswith((".beta", ".eps")):
+                buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
+
+    context = _identity(cold_model, source)
+    runtime = _runtime(tmp_path / "flux2-klein-store")
+    miss = runtime.resolve(context.identity)
+    assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+    producer = FinalLayoutBF16Producer(
+        context,
+        cold_model,
+        (("transformer", cold_transformer),),
+        max_shard_bytes=64,
+    )
+    publication = runtime.publish_after_load(context.identity, producer=producer)
+    assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
+
+    warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
+    warm_model = _Flux2KleinPipeline(warm_transformer)
+    warm_context = _identity(warm_model, source)
+    assert warm_context.identity == context.identity
+    hit = runtime.resolve(warm_context.identity)
+    assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
+    assert hit.lease is not None
+
+    plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
+    assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
+    plan.commit()
+    warm_transformer.validate_restored_host_weights()
+
+    cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
+    warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
+    for name in warm_context.tensor_names:
+        assert torch.equal(warm_tensors[name], cold_tensors[name])
+        assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
+    assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
+
+    del warm_tensors, warm_model, warm_transformer, plan
+    gc.collect()
+    hit.lease.close()
+    assert hit.lease.closed
+    assert isinstance(runtime.store, FilesystemHostWeightStore)
+    assert runtime.store.cleanup(context.identity) is None

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -35,7 +35,10 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadHook,
     PinnedResidentLayerGroup,
 )
-from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError
+from vllm_omni.diffusion.offloader.host_registration import (
+    HostRegistrationCleanupError,
+    HostRegistrationError,
+)
 from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
@@ -807,7 +810,10 @@ def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_ru
     assert torch.equal(block.weight, expected)
 
 
-def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch):
+def test_hwr_registration_uses_transport_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    patched_offload_runtime,
+):
     dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
     plan, _, lease = _fake_hwr_plan()
     backend = DistributedLayerwiseOffloadBackend(
@@ -1125,6 +1131,51 @@ class TestMmapWeightLoading:
         backend.disable()
         assert lease.closed
 
+    def test_hwr_registration_failure_falls_back_to_bounded_staging(
+        self,
+        monkeypatch,
+        patched_offload_runtime,
+    ):
+        plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dp_size=1,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        pipeline = _HWRPipeline()
+        original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
+
+        def fail_registration(*args, **kwargs):
+            del args, kwargs
+            raise HostRegistrationError("registration unavailable in CPU test")
+
+        def allocate_unpinned_staging(hooks, resident_group=None):
+            # Registration selection still observes pin_cpu_memory=True. The
+            # CPU-only test avoids asking the host for CUDA-pinned allocations.
+            for hook in hooks:
+                hook.pin_memory = False
+            return original_staging_allocator(hooks, resident_group)
+
+        monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
+        monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
+
+        backend.enable(pipeline)
+
+        assert carrier.taken
+        assert backend._using_rank_local_mmap
+        assert not backend._using_registered_mmap
+        assert backend._host_registration is None
+        assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
+        assert not lease.closed
+
+        backend.disable()
+        assert lease.closed
+
     def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
         self,
         monkeypatch,
@@ -2132,8 +2183,9 @@ class TestDynamicSlotTracking:
                 prev_idx = (i - 1) % num_blocks
 
                 # Dynamic slot tracking: read from prev's prefetched slot
-                if prefetched_slots[prev_idx] is not None:
+                prefetched_slot = prefetched_slots[prev_idx]
-                    current_slots[i] = prefetched_slots[prev_idx]
+                if prefetched_slot is not None:
+                    current_slots[i] = prefetched_slot
 
                 read_slot = current_slots[i]
 

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -49,12 +49,23 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
+from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 
 logger = init_logger(__name__)
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 
+_FLUX2_STACKED_PARAMS_MAPPING = (
+    (".to_qkv.", ".to_q.", "q"),
+    (".to_qkv.", ".to_k.", "k"),
+    (".to_qkv.", ".to_v.", "v"),
+    (".add_kv_proj", ".add_q_proj", "q"),
+    (".add_kv_proj", ".add_k_proj", "k"),
+    (".add_kv_proj", ".add_v_proj", "v"),
+)
+
+
 def _join_prefix(prefix: str, name: str) -> str:
     return f"{prefix}.{name}" if prefix else name
 
@@ -770,6 +781,14 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    # Constructor state plus complete final-layout parameters and persistent
+    # buffers is sufficient to reconstruct a ready base-model transformer.
+    # Changes to the packed mapping or restore invariants require a version bump.
+    host_weight_restore_contract = FinalLayoutModelContract(
+        implementation_id="flux2-klein-dit",
+        version="1",
+    )
+
     _cache_dit_adapter_config = CacheDiTAdapterConfig(
         block_forward_patterns={
             "transformer_blocks": ForwardPattern.Pattern_1,
@@ -817,6 +836,9 @@ class Flux2Transformer2DModel(nn.Module):
         quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
+        # Warm HWR restoration skips load_weights(), so loader consumers such
+        # as LoRA discovery must not depend on that method to create this state.
+        self.stacked_params_mapping = list(_FLUX2_STACKED_PARAMS_MAPPING)
         self.out_channels = out_channels or in_channels
         self.inner_dim = num_attention_heads * attention_head_dim
         self.config = SimpleNamespace(
@@ -1004,18 +1026,78 @@ class Flux2Transformer2DModel(nn.Module):
             return (output,)
         return Transformer2DModelOutput(sample=output)
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def validate_restored_host_weights(self) -> None:
-        stacked_params_mapping = [
+        """Validate the tensor-complete BF16 base-model restore contract."""
-            (".to_qkv.", ".to_q.", "q"),
+        if not isinstance(self.stacked_params_mapping, (list, tuple)) or (
-            (".to_qkv.", ".to_k.", "k"),
+            tuple(self.stacked_params_mapping) != _FLUX2_STACKED_PARAMS_MAPPING
-            (".to_qkv.", ".to_v.", "v"),
+        ):
-            (".add_kv_proj", ".add_q_proj", "q"),
+            raise ValueError("FLUX.2-klein packed QKV mapping is not constructor-stable")
-            (".add_kv_proj", ".add_k_proj", "k"),
-            (".add_kv_proj", ".add_v_proj", "v"),
-        ]
-        # Expose packed shard mappings for LoRA handling of fused projections.
-        self.stacked_params_mapping = stacked_params_mapping
 
+        stack_specs = (
+            ("transformer_blocks", self.config.num_layers, Flux2TransformerBlock),
+            ("single_transformer_blocks", self.config.num_single_layers, Flux2SingleTransformerBlock),
+        )
+        for name, expected_count, block_type in stack_specs:
+            blocks = getattr(self, name, None)
+            if not isinstance(blocks, nn.ModuleList) or len(blocks) != expected_count or expected_count <= 0:
+                raise ValueError(
+                    f"FLUX.2-klein {name} is incomplete: expected {expected_count} blocks, "
+                    f"got {len(blocks) if isinstance(blocks, nn.ModuleList) else 'missing'}"
+                )
+            if any(not isinstance(block, block_type) for block in blocks):
+                raise ValueError(f"FLUX.2-klein {name} contains an unexpected block implementation")
+
+        required_modules = (
+            "pos_embed",
+            "rope_prepare",
+            "time_guidance_embed",
+            "double_stream_modulation_img",
+            "double_stream_modulation_txt",
+            "single_stream_modulation",
+            "x_embedder",
+            "context_embedder",
+            "norm_out",
+            "proj_out",
+        )
+        for name in required_modules:
+            if not isinstance(getattr(self, name, None), nn.Module):
+                raise ValueError(f"FLUX.2-klein required module {name!r} is missing")
+
+        parameter_count = 0
+        stack_parameter_coverage = {name: False for name, _, _ in stack_specs}
+        for name, parameter in self.named_parameters():
+            parameter_count += 1
+            for stack_name in stack_parameter_coverage:
+                if name.startswith(f"{stack_name}."):
+                    stack_parameter_coverage[stack_name] = True
+            if parameter.is_meta or parameter.device.type != "cpu":
+                raise ValueError(f"FLUX.2-klein parameter {name!r} is not materialized on CPU")
+            if parameter.dtype is not torch.bfloat16:
+                raise ValueError(f"FLUX.2-klein parameter {name!r} must stay bf16, got {parameter.dtype}")
+            if parameter.layout is not torch.strided or not parameter.is_contiguous():
+                raise ValueError(f"FLUX.2-klein parameter {name!r} must use contiguous strided layout")
+            if parameter.numel() == 0:
+                raise ValueError(f"FLUX.2-klein parameter {name!r} is empty")
+        if parameter_count == 0 or not all(stack_parameter_coverage.values()):
+            raise ValueError("FLUX.2-klein restore did not materialize both transformer block stacks")
+
+        for name, buffer in self.named_buffers():
+            parent_path, _, leaf_name = name.rpartition(".")
+            owner = self.get_submodule(parent_path)
+            persistent = leaf_name not in owner._non_persistent_buffers_set
+            loader_buffer = name.endswith((".beta", ".eps"))
+            if loader_buffer and not persistent:
+                raise ValueError(f"FLUX.2-klein loader buffer {name!r} must be persistent")
+            if not persistent:
+                continue
+            if buffer.is_meta or buffer.device.type != "cpu":
+                raise ValueError(f"FLUX.2-klein persistent buffer {name!r} is not materialized on CPU")
+            if buffer.layout is not torch.strided or not buffer.is_contiguous():
+                raise ValueError(f"FLUX.2-klein persistent buffer {name!r} must use contiguous strided layout")
+            if loader_buffer and (buffer.dtype is not torch.bfloat16 or buffer.numel() != 1):
+                raise ValueError(f"FLUX.2-klein loader buffer {name!r} must be a scalar bf16 tensor")
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         params_dict = dict(self.named_parameters())
 
         for name, buffer in self.named_buffers():
@@ -1026,7 +1108,7 @@ class Flux2Transformer2DModel(nn.Module):
         for name, loaded_weight in weights:
             original_name = name
             mapped = False
-            for param_name, weight_name, shard_id in stacked_params_mapping:
+            for param_name, weight_name, shard_id in self.stacked_params_mapping:
                 if weight_name not in original_name:
                     continue
                 name = original_name.replace(weight_name, param_name)