[Diffusion] Add FLUX.2-klein BF16 HWR contract (#6651)

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.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 initial diffusion consumer contract additionally defines typed,
+The implementation provides contracts and a CPU local-filesystem store. The
+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
-library-only: no loader or DLO path selects, publishes, restores, or transports
-that artifact yet.
+concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
+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;
-- lease handoff to DLO or another transport consumer;
-- FP8, quantized, merged-adaptation, or additional model producers;
-- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
+- default-on activation or consumers outside no-AllGather DLO;
+- online FP8, quantized, merged-adaptation, or additional model producers;
+- HWR interaction with DLO AllGather;
 - 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
-and TP2 rank identities plus SP layout identities; online quantization, HSDP,
-LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode
+contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
+contract covers both transformer block stacks, constructor-stable packed QKV
+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 \