From 4e39c68c6b19037cc6e2f7d4fbe8658014719297 Mon Sep 17 00:00:00 2001 From: yiming-l21 Date: Wed, 12 Aug 2026 12:02:34 +0800 Subject: [PATCH] feat: add Flux2 Klein runtime support --- README.md | 19 +- docs/models.md | 1 + src/core/optimization/cache/cache_types.cpp | 5 + .../cache/model/dit_model_cache_contract.cpp | 6 + src/core/runtime/model_loader.cpp | 76 +++++- src/dit_models/components/common/rope.hpp | 10 +- .../components/text_encoders/llm.hpp | 2 +- src/dit_models/models/flux.hpp | 32 ++- src/dit_models/pipelines/flux_pipeline.cpp | 199 +++++++++++++-- third_party/ggml | 2 +- ...ench_diffusers_flux_transformer_ulysses.py | 230 +++++++++++++++++ tools/bench_diffusers_flux_ulysses.py | 120 +++++++++ tools/bench_diffusers_wan_transformer.py | 233 ++++++++++++++++++ 13 files changed, 884 insertions(+), 51 deletions(-) create mode 100644 tools/bench_diffusers_flux_transformer_ulysses.py create mode 100644 tools/bench_diffusers_flux_ulysses.py create mode 100644 tools/bench_diffusers_wan_transformer.py diff --git a/README.md b/README.md index 29cacc8..b7eb988 100644 --- a/README.md +++ b/README.md @@ -18,9 +18,10 @@ runtime for image generation, image editing, and video generation across local, edge, and resource-constrained deployment environments. -It supports major DiT model families including **FLUX.1, Stable Diffusion 3/3.5, -Qwen-Image, and Wan**, with explicit control over model loading, memory usage, -graph execution, quantization, device placement, and backend selection. +It supports major DiT model families including **FLUX.1, FLUX.2 [klein] 4B, +Stable Diffusion 3/3.5, Qwen-Image, and Wan**, with explicit control over model +loading, memory usage, graph execution, quantization, device placement, and +backend selection. ## Features @@ -32,7 +33,8 @@ graph execution, quantization, device placement, and backend selection. - **Unified across tasks and model families** - **Text-to-image**, **image editing**, and **video generation** in one runtime - - SD3/SD3.5, FLUX.1, FLUX.1-Kontext, Qwen-Image, Qwen-Image-Edit, and Wan 2.1 + - SD3/SD3.5, FLUX.1, FLUX.1-Kontext, FLUX.2 [klein] 4B, Qwen-Image, + Qwen-Image-Edit, and Wan 2.1 - **Few-step distilled models** auto-detected — Turbo / Lightning / schnell default to a **4–8 step** schedule - Shared **C API, CLI, HTTP server, and Python** interfaces across every family @@ -63,10 +65,10 @@ graph execution, quantization, device placement, and backend selection. ## Supported Models -This release focuses on the model families below, each with a base -checkpoint and a **few-step distilled variant**. Some source files contain -experimental model scaffolding beyond this table; those are not part of the -current support commitment unless documented in +This release focuses on the model families below. Most ship as a base +checkpoint plus a **few-step distilled variant**; some source files contain +experimental model scaffolding beyond this table, and those are not part of +the current support commitment unless documented in [Supported Models](docs/models.md). | Model family | Task | Base checkpoint | Distilled variant (few-step) | Status | @@ -74,6 +76,7 @@ current support commitment unless documented in | **SD3 / SD3.5** | Text-to-image | `stabilityai/stable-diffusion-3-medium` | SD3.5-medium-turbo | Supported | | **FLUX.1** | Text-to-image | `black-forest-labs/FLUX.1-dev` | FLUX.1-schnell | Supported | | **FLUX.1-Kontext** | Image editing / reference-guided | `black-forest-labs/FLUX.1-Kontext-dev` | Kontext Lightning | Supported | +| **FLUX.2 [klein] 4B** | Text-to-image / image editing | `black-forest-labs/FLUX.2-klein-4B` | Native FLUX.2 [klein] 4B checkpoint | Supported | | **Qwen-Image** | Text-to-image | `Qwen/Qwen-Image` | Qwen-Image Lightning *(LoRA)* | Supported | | **Qwen-Image-Edit** | Image editing | `Qwen/Qwen-Image-Edit` | Qwen-Image-Edit Lightning *(LoRA)* | Supported | | **Wan 2.1** | Video generation | `Wan-AI/Wan2.1-T2V-1.3B` (and 14B) | Wan2.1-T2V-1.3B Distill | Supported (Vulkan still optimizing) | diff --git a/docs/models.md b/docs/models.md index 3b54599..7a98010 100644 --- a/docs/models.md +++ b/docs/models.md @@ -15,6 +15,7 @@ support commitment. | SD3 / SD3.5 | Text-to-image | [`stabilityai/stable-diffusion-3-medium-diffusers`](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) (and SD3.5 siblings) | Diffusers-style directory or component weights | CUDA first, CPU/Vulkan functional, Metal experimental | Supported | | FLUX.1 | Text-to-image | [`black-forest-labs/FLUX.1-dev`](https://huggingface.co/black-forest-labs/FLUX.1-dev) | Diffusers-style directory, top-level FLUX safetensors, or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported | | FLUX.1-Kontext | Image editing / reference-guided generation | [`black-forest-labs/FLUX.1-Kontext-dev`](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported | +| FLUX.2 [klein] 4B | Text-to-image / image editing | [`black-forest-labs/FLUX.2-klein-4B`](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported | | Qwen-Image | Text-to-image | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported | | Qwen-Image-Edit | Image editing | [`Qwen/Qwen-Image-Edit`](https://huggingface.co/Qwen/Qwen-Image-Edit) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported | | Wan 2.1 | Video generation | [`Wan-AI/Wan2.1-T2V-1.3B-Diffusers`](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) (and Wan2.1-T2V-14B-Diffusers) | Diffusers-style directory or components | CUDA first, CPU functional for validation, Metal/Vulkan experimental | Supported (Vulkan still optimizing) | diff --git a/src/core/optimization/cache/cache_types.cpp b/src/core/optimization/cache/cache_types.cpp index 02fec51..2fed843 100644 --- a/src/core/optimization/cache/cache_types.cpp +++ b/src/core/optimization/cache/cache_types.cpp @@ -34,6 +34,11 @@ CacheModelSpec cache_model_spec_for_version(SDVersion version) { CacheModelSpec spec; spec.version = version; + if (version == VERSION_FLUX2_KLEIN) { + spec.model_name = "flux2-klein"; + spec.block_count = 5 + 20; // double_blocks + single_blocks + return spec; + } if (ed_version_is_flux(version) || ed_version_is_flux2(version)) { spec.model_name = "flux"; spec.block_count = 19 + 38; // double_blocks + single_blocks diff --git a/src/core/optimization/cache/model/dit_model_cache_contract.cpp b/src/core/optimization/cache/model/dit_model_cache_contract.cpp index 15ed7fb..3539bfd 100644 --- a/src/core/optimization/cache/model/dit_model_cache_contract.cpp +++ b/src/core/optimization/cache/model/dit_model_cache_contract.cpp @@ -18,6 +18,12 @@ struct FamilyInfo { FamilyInfo family_info_for(SDVersion version) { FamilyInfo info; + if (version == VERSION_FLUX2_KLEIN) { + info.family = ModelFamily::Flux; + info.num_blocks = 5 + 20; // double_blocks + single_blocks + info.dual_stream = true; + return info; + } if (ed_version_is_flux(version) || ed_version_is_flux2(version)) { info.family = ModelFamily::Flux; info.num_blocks = 19 + 38; // double_blocks + single_blocks diff --git a/src/core/runtime/model_loader.cpp b/src/core/runtime/model_loader.cpp index 5569ddf..5f80c45 100644 --- a/src/core/runtime/model_loader.cpp +++ b/src/core/runtime/model_loader.cpp @@ -359,6 +359,9 @@ static SDVersion infer_transformer_file_version(const std::string& file_path) { if (contains(klass, "Kontext")) { return VERSION_FLUX_KONTEXT; } + if (contains(klass, "Flux2")) { + return contains(klass, "Klein") ? VERSION_FLUX2_KLEIN : VERSION_FLUX2; + } if (contains(klass, "Flux")) { return VERSION_FLUX; } @@ -377,6 +380,10 @@ static bool is_flux1_family_version(SDVersion version) { return version == VERSION_FLUX || version == VERSION_FLUX_KONTEXT; } +static bool is_flux2_family_version(SDVersion version) { + return ed_version_is_flux2(version); +} + static std::string resolve_flux_transformer_component_path(const std::string& file_path) { std::string normalized = file_path; std::replace(normalized.begin(), normalized.end(), '\\', '/'); @@ -450,6 +457,9 @@ static SDVersion infer_diffusers_version(const std::string& dir_path) { if (contains(klass, "Kontext")) { return VERSION_FLUX_KONTEXT; } + if (contains(klass, "Flux2")) { + return contains(klass, "Klein") ? VERSION_FLUX2_KLEIN : VERSION_FLUX2; + } if (contains(klass, "Flux")) { return VERSION_FLUX; } @@ -478,6 +488,9 @@ static SDVersion infer_diffusers_version(const std::string& dir_path) { if (contains(klass, "Kontext")) { return VERSION_FLUX_KONTEXT; } + if (contains(klass, "Flux2")) { + return contains(klass, "Klein") ? VERSION_FLUX2_KLEIN : VERSION_FLUX2; + } if (contains(klass, "Flux")) { return VERSION_FLUX; } @@ -539,6 +552,29 @@ static std::vector component_weight_candidates(const std::string& c return candidates; } +static std::string find_top_level_safetensors_file(const std::string& dir_path) { + std::error_code ec; + if (!fs::is_directory(dir_path, ec)) { + return {}; + } + + for (const auto& entry : fs::directory_iterator(dir_path, ec)) { + if (ec) { + break; + } + if (!entry.is_regular_file()) { + continue; + } + + const std::string path = entry.path().string(); + if (has_suffix(path, ".safetensors")) { + return path; + } + } + + return {}; +} + static uint16_t f8_e4m3_to_f16(uint8_t f8) { const uint32_t exponent_bias = 7; if (f8 == 0xff) { @@ -1281,7 +1317,7 @@ bool ModelLoader::init_from_diffusers_directory(const std::string& dir_path, con } const std::string component_dir = path_join(dir_path, component.dir); if (!is_directory(component_dir)) { - if (component.required_for_flux && is_flux1_family_version(version_)) { + if (component.required_for_flux && (is_flux1_family_version(version_) || is_flux2_family_version(version_))) { LOG_WARN("diffusers component '%s' not found", component.dir); } continue; @@ -1297,6 +1333,15 @@ bool ModelLoader::init_from_diffusers_directory(const std::string& dir_path, con continue; } } + if (is_flux2_family_version(version_)) { + if (std::strcmp(component.dir, "text_encoder") == 0) { + component_prefix = "text_encoders.llm."; + } else if (std::strcmp(component.dir, "text_encoder_2") == 0 || + std::strcmp(component.dir, "text_encoder_3") == 0 || + std::strcmp(component.dir, "unet") == 0) { + continue; + } + } if (ed_version_is_qwen_image(version_) || ed_version_is_qwen_image_edit(version_)) { if (std::strcmp(component.dir, "text_encoder") == 0) { component_prefix = "text_encoders.llm."; @@ -1308,6 +1353,20 @@ bool ModelLoader::init_from_diffusers_directory(const std::string& dir_path, con } bool loaded = false; std::set tried; + + if (is_flux2_family_version(version_) && std::strcmp(component.dir, "transformer") == 0) { + const std::string top_level_flux = find_top_level_safetensors_file(dir_path); + if (!top_level_flux.empty()) { + const size_t before = tensor_storage_map_.size(); + loaded = init_from_safetensors_file(top_level_flux, component_prefix); + if (loaded) { + LOG_INFO("loaded diffusers component '%s' from top-level Flux2 weights '%s' (%zu tensors)", + component.dir, + top_level_flux.c_str(), + tensor_storage_map_.size() - before); + } + } + } if (is_flux1_family_version(version_) && std::strcmp(component.dir, "transformer") == 0) { const std::vector top_level_flux_weights = { path_join(dir_path, "flux1-kontext-dev.safetensors"), @@ -1390,6 +1449,8 @@ SDVersion ModelLoader::get_ld_version() { bool has_transformer_blocks = false; bool has_unet = false; bool has_second_text_encoder = false; + bool has_flux2 = false; + bool has_single_block_47 = false; TensorStorage input_block_weight; TensorStorage token_embedding_weight; @@ -1417,9 +1478,16 @@ SDVersion ModelLoader::get_ld_version() { if (contains(name, "model.diffusion_model.double_blocks.") || contains(name, "transformer.double_blocks.")) { has_flux_double = true; } + if (contains(name, "model.diffusion_model.double_stream_modulation_img.lin.weight") || + contains(name, "model.diffusion_model.double_stream_modulation_img.linear.weight")) { + has_flux2 = true; + } if (contains(name, "single_transformer_blocks.") || contains(name, "single_blocks.")) { has_flux_single = true; } + if (contains(name, "single_blocks.47.linear1.weight")) { + has_single_block_47 = true; + } if (contains(name, "transformer.transformer_blocks.") || contains(name, "transformer_blocks.")) { has_transformer_blocks = true; } @@ -1437,6 +1505,12 @@ SDVersion ModelLoader::get_ld_version() { } } + if (has_flux2) { + if (has_single_block_47) { + return VERSION_FLUX2; + } + return VERSION_FLUX2_KLEIN; + } if (has_flux_double || (has_transformer_blocks && has_flux_single)) { if (input_block_weight.ne[0] == 384) { return VERSION_FLUX_FILL; diff --git a/src/dit_models/components/common/rope.hpp b/src/dit_models/components/common/rope.hpp index 5d3d150..2358bff 100644 --- a/src/dit_models/components/common/rope.hpp +++ b/src/dit_models/components/common/rope.hpp @@ -785,6 +785,7 @@ namespace Rope { ggml_tensor* mask, float kv_scale = 1.0f, bool rope_interleaved = true, + bool use_fused_rope = true, bool k_rope_f16_for_flash = false, bool qk_rope_bf16_roundtrip = false) { // q,k,v: [N, L, n_head, d_head] @@ -796,16 +797,17 @@ namespace Rope { ctx->flash_attn_enabled && k->ne[2] % 256 != 0; - q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved, ctx->backend); // [N*n_head, L, d_head] - if (k_rope_f16_for_flash && ctx->flash_attn_enabled && !will_pad_kv_for_flash_attn) { + ggml_backend_t rope_backend = use_fused_rope ? ctx->backend : nullptr; + q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved, rope_backend); // [N*n_head, L, d_head] + if (k_rope_f16_for_flash && use_fused_rope && ctx->flash_attn_enabled && !will_pad_kv_for_flash_attn) { ggml_tensor* k_f16 = edgedit::ggml_ext::apply_rope_f16(ctx->ggml_ctx, k, pe, rope_interleaved); if (k_f16 != nullptr && ggml_backend_supports_op(ctx->backend, k_f16)) { k = k_f16; // flash attention consumes K as F16, so avoid a separate cast node } else { - k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, ctx->backend); + k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, rope_backend); } } else { - k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, ctx->backend); // [N*n_head, L, d_head] + k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, rope_backend); // [N*n_head, L, d_head] } if (qk_rope_bf16_roundtrip) { q = ggml_cast(ctx->ggml_ctx, ggml_cast(ctx->ggml_ctx, q, GGML_TYPE_BF16), GGML_TYPE_F32); diff --git a/src/dit_models/components/text_encoders/llm.hpp b/src/dit_models/components/text_encoders/llm.hpp index b110f3a..14c6125 100644 --- a/src/dit_models/components/text_encoders/llm.hpp +++ b/src/dit_models/components/text_encoders/llm.hpp @@ -1612,7 +1612,7 @@ namespace LLM { params.rms_norm_eps = 1e-5f; } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) { params.head_dim = 128; - params.num_heads = 64; + params.num_heads = 32; params.num_kv_heads = 8; params.qkv_bias = false; params.qk_norm = true; diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index a0d563b..3e66a95 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -1637,6 +1637,7 @@ namespace Flux { struct DoubleStreamBlock : public GGMLBlock { bool prune_mod; int idx = 0; + bool use_fused_rope = true; public: DoubleStreamBlock(int64_t hidden_size, @@ -1649,8 +1650,9 @@ namespace Flux { bool mlp_proj_bias = true, bool use_yak_mlp = false, bool use_mlp_silu_act = false, - bool preserve_activation_dtype = false) - : idx(idx), prune_mod(prune_mod) { + bool preserve_activation_dtype = false, + bool use_fused_rope = true) + : idx(idx), prune_mod(prune_mod), use_fused_rope(use_fused_rope) { int64_t mlp_hidden_dim = static_cast(hidden_size * mlp_ratio); if (!prune_mod && !share_modulation) { @@ -1846,7 +1848,9 @@ namespace Flux { flux_align_debug_capture("block0.attn.joint_v", v); } - auto attn = Rope::attention(ctx, q, k, v, pe, mask, 1.0f, true, true); // [N, n_txt_token + n_img_token, n_head*d_head] + auto attn = use_fused_rope + ? Rope::attention(ctx, q, k, v, pe, mask, 1.0f, true, true) + : Rope::attention(ctx, q, k, v, pe, mask, 1.0f, true, false); // [N, n_txt_token + n_img_token, n_head*d_head] attn = flux_cast_activation(ctx->ggml_ctx, attn, img->type); if (idx == 0) { flux_align_debug_capture("block0.attn.joint_out", attn); @@ -2907,6 +2911,7 @@ namespace Flux { int idx = 0; bool use_yak_mlp; bool use_mlp_silu_act; + bool use_fused_rope = true; int64_t mlp_mult_factor; public: @@ -2920,8 +2925,9 @@ namespace Flux { bool mlp_proj_bias = true, bool use_yak_mlp = false, bool use_mlp_silu_act = false, - bool preserve_activation_dtype = false) - : hidden_size(hidden_size), num_heads(num_heads), idx(idx), prune_mod(prune_mod), use_yak_mlp(use_yak_mlp), use_mlp_silu_act(use_mlp_silu_act) { + bool preserve_activation_dtype = false, + bool use_fused_rope = true) + : hidden_size(hidden_size), num_heads(num_heads), idx(idx), prune_mod(prune_mod), use_yak_mlp(use_yak_mlp), use_mlp_silu_act(use_mlp_silu_act), use_fused_rope(use_fused_rope) { int64_t head_dim = hidden_size / num_heads; float scale = qk_scale; if (scale <= 0.f) { @@ -3007,7 +3013,9 @@ namespace Flux { q = norm->query_norm(ctx, q); k = norm->key_norm(ctx, k); - auto attn = Rope::attention(ctx, q, k, v, pe, mask, 1.0f, true, true); // [N, n_token, hidden_size] + auto attn = use_fused_rope + ? Rope::attention(ctx, q, k, v, pe, mask, 1.0f, true, true) + : Rope::attention(ctx, q, k, v, pe, mask, 1.0f, true, false); // [N, n_token, hidden_size] attn = flux_cast_activation(ctx->ggml_ctx, attn, x->type); auto mlp = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, mlp_hidden_dim * mlp_mult_factor, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], hidden_size * 3 * qkv_mlp->nb[0]); @@ -3730,6 +3738,7 @@ namespace Flux { bool semantic_txt_norm = false; bool use_yak_mlp = false; bool use_mlp_silu_act = false; + bool use_fused_rope = true; float ref_index_scale = 1.f; ggml_type activation_dtype = GGML_TYPE_F32; ChromaRadianceParams chroma_radiance_params; @@ -3802,7 +3811,8 @@ namespace Flux { !params.disable_bias, params.use_yak_mlp, params.use_mlp_silu_act, - preserve_activation_dtype); + preserve_activation_dtype, + params.use_fused_rope); } for (int i = 0; i < params.depth_single_blocks; i++) { @@ -3816,7 +3826,8 @@ namespace Flux { !params.disable_bias, params.use_yak_mlp, params.use_mlp_silu_act, - preserve_activation_dtype); + preserve_activation_dtype, + params.use_fused_rope); } if (params.version == VERSION_CHROMA_RADIANCE) { @@ -4792,6 +4803,7 @@ namespace Flux { flux_params.share_modulation = true; flux_params.ref_index_scale = 10.f; flux_params.use_mlp_silu_act = true; + flux_params.use_fused_rope = false; } uint32_t diffusion_tensor_count = 0; uint32_t diffusion_bf16_count = 0; @@ -4866,6 +4878,8 @@ namespace Flux { flux_params.activation_dtype = flux_env_flag_enabled("ED_FLUX_BF16_ACTIVATIONS") ? GGML_TYPE_BF16 : GGML_TYPE_F32; + } else if (ed_version_is_flux2(version)) { + flux_params.activation_dtype = GGML_TYPE_F32; } LOG_INFO("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 @@ -5121,7 +5135,7 @@ namespace Flux { } // Dummy shape-only inputs matching the real compute() call at flux_pipeline // (x, timesteps, context, {}, y, guidance). Data is never read during measure. - sd::Tensor x = sd::zeros({latent_w, latent_h, 16, 1}); + sd::Tensor x = sd::zeros({latent_w, latent_h, static_cast(flux_params.in_channels), 1}); sd::Tensor timesteps = sd::zeros({1}); // T5 context: [context_in_dim, tokens, 1]. Use 512 tokens (flux T5 max) so the // attention activation is measured at/above the real sequence length. diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index 5d3d35a..4d03f45 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -52,6 +52,32 @@ static std::string format_type_counts(const std::map& type_ return ss.str(); } +static bool flux_pipeline_uses_llm_conditioner(SDVersion version) { + return ed_version_is_flux2(version); +} + +static bool flux_pipeline_uses_flash_attention(SDVersion version, bool runtime_flash_attention) { + (void)version; + return runtime_flash_attention; +} + +static const char* flux_pipeline_required_weights(SDVersion version) { + if (flux_pipeline_uses_llm_conditioner(version)) { + return "transformer, text_encoder/LLM, and VAE"; + } + return "transformer, CLIP-L, T5XXL, and VAE"; +} + +static int flux_pipeline_latent_channels(SDVersion version) { + if (ed_version_uses_flux2_vae(version)) { + return 128; + } + if (version == VERSION_CHROMA_RADIANCE) { + return 3; + } + return 16; +} + template static std::string format_tensor_shape(const sd::Tensor& tensor) { if (tensor.empty()) { @@ -205,6 +231,49 @@ static std::vector ed_flux_discrete_sigmas(int steps, float shift) { return result; } +static float ed_flux2_empirical_mu(int image_seq_len, int steps) { + const float a1 = 8.73809524e-05f; + const float b1 = 1.89833333f; + const float a2 = 0.00016927f; + const float b2 = 0.45666666f; + + if (image_seq_len > 4300) { + return a2 * static_cast(image_seq_len) + b2; + } + + const float m_200 = a2 * static_cast(image_seq_len) + b2; + const float m_10 = a1 * static_cast(image_seq_len) + b1; + const float a = (m_200 - m_10) / 190.0f; + const float b = m_200 - 200.0f * a; + return a * static_cast(steps) + b; +} + +static std::vector ed_flux2_sigmas(int steps, int image_seq_len, float* out_mu) { + std::vector result; + if (steps <= 0) { + return result; + } + + const float mu = ed_flux2_empirical_mu(image_seq_len, steps); + if (out_mu != nullptr) { + *out_mu = mu; + } + + result.reserve(static_cast(steps) + 1); + for (int i = 0; i <= steps; ++i) { + const float t = 1.0f - static_cast(i) / static_cast(steps); + if (t <= 0.0f) { + result.push_back(0.0f); + } else if (t >= 1.0f) { + result.push_back(1.0f); + } else { + result.push_back(ed_flux_time_shift(mu, 1.0f, t)); + } + } + result[static_cast(steps)] = 0.0f; + return result; +} + static ed_status_t ed_tensor_to_image(const sd::Tensor& tensor, ed_image_t* image) { if (image == nullptr || tensor.empty()) { return ED_STATUS_INVALID_ARGUMENT; @@ -412,14 +481,26 @@ bool FluxPipeline::validate(std::string* error) const { } return false; } - if (!has_component("clip_l")) { - LOG_WARN("Flux manifest has no CLIP-L text encoder tensors; this is OK for transformer-only files"); - } - if (!has_component("t5xxl")) { - LOG_WARN("Flux manifest has no T5XXL text encoder tensors; this is OK for transformer-only files"); - } - if (!has_component("vae")) { - LOG_WARN("Flux manifest has no VAE tensors; this is OK for transformer-only files"); + if (ed_version_is_flux2(version_)) { + if (!has_component("text_encoder")) { + if (error != nullptr) { + *error = "Flux2 model is missing text_encoder/LLM tensors"; + } + return false; + } + if (!has_component("vae")) { + LOG_WARN("Flux2 manifest has no VAE tensors; this is OK for transformer-only files"); + } + } else { + if (!has_component("clip_l")) { + LOG_WARN("Flux manifest has no CLIP-L text encoder tensors; this is OK for transformer-only files"); + } + if (!has_component("t5xxl")) { + LOG_WARN("Flux manifest has no T5XXL text encoder tensors; this is OK for transformer-only files"); + } + if (!has_component("vae")) { + LOG_WARN("Flux manifest has no VAE tensors; this is OK for transformer-only files"); + } } } @@ -452,7 +533,7 @@ bool FluxPipeline::initialize_flux_transformer_spec(const ModelLoader& loader, version_, false)); if (runtime_ != nullptr) { - const bool diffusion_flash = runtime_->flash_attention(); + const bool diffusion_flash = flux_pipeline_uses_flash_attention(version_, runtime_->flash_attention()); flux_runner_->set_max_graph_vram_bytes(runtime_->max_graph_vram_bytes()); flux_runner_->set_flash_attention_enabled(diffusion_flash); @@ -597,7 +678,42 @@ bool FluxPipeline::prepare_flux_runtime_weights(const ModelLoader& loader, flux_runner_->get_param_tensors(registry.tensors(), "model.diffusion_model"); - if (has_component("clip_l") || has_component("t5xxl")) { + const bool flux_flash_attention = + runtime_ != nullptr ? flux_pipeline_uses_flash_attention(version_, runtime_->flash_attention()) : false; + + if (flux_pipeline_uses_llm_conditioner(version_)) { + if (!has_component("text_encoder")) { + if (error != nullptr) { + *error = "Flux2 model is missing text_encoder tensors"; + } + return false; + } + if (text_backend == nullptr) { + if (error != nullptr) { + *error = "FluxPipeline requires a non-null text encoder backend from ModelRuntime"; + } + return false; + } + conditioner_backend_ = text_backend; + + conditioner_ = std::make_shared(conditioner_backend_, + te_offload, + loader.get_tensor_storage_map(), + version_, + "", + false); + + conditioner_->alloc_params_buffer(); + conditioner_->get_param_tensors(registry.tensors()); + registry.ignore_prefix("text_encoders.llm.lm_head."); + registry.ignore_prefix("text_encoders.llm.output.weight"); + registry.ignore_prefix("text_encoders.llm.visual."); + // TE params buffer now allocated: real weight size is known. Set a TE-specific + // segment budget so an offloaded text encoder segments instead of staging whole. + conditioner_->set_max_graph_vram_bytes( + runtime_->text_encoder_segment_budget(conditioner_->get_params_buffer_size())); + conditioner_->set_flash_attention_enabled(flux_flash_attention); + } else if (has_component("clip_l") || has_component("t5xxl")) { if (text_backend == nullptr) { if (error != nullptr) { *error = "FluxPipeline requires a non-null text encoder backend from ModelRuntime"; @@ -617,6 +733,7 @@ bool FluxPipeline::prepare_flux_runtime_weights(const ModelLoader& loader, // whole. No-op for a resident TE (returns the global budget). conditioner_->set_max_graph_vram_bytes( runtime_->text_encoder_segment_budget(conditioner_->get_params_buffer_size())); + conditioner_->set_flash_attention_enabled(flux_flash_attention); } if (has_component("vae")) { @@ -735,7 +852,8 @@ ed_status_t FluxPipeline::generate_image(const ed_image_generation_params_t* par } if (!can_generate_image()) { if (error != nullptr) { - *error = "current Flux pipeline needs transformer, CLIP-L, T5XXL, and VAE weights"; + *error = std::string("current Flux pipeline needs ") + + flux_pipeline_required_weights(version_) + " weights"; } return ED_STATUS_UNSUPPORTED; } @@ -819,7 +937,8 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params } if (!can_generate_image()) { if (error != nullptr) { - *error = "full Flux runtime is not loaded; need transformer, CLIP-L, T5XXL, and VAE weights"; + *error = std::string("full Flux runtime is not loaded; need ") + + flux_pipeline_required_weights(version_) + " weights"; } return false; } @@ -882,10 +1001,9 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params format_tensor_shape(condition.c_vector).c_str()); const int steps = resolve_steps(params->sample.steps); + const bool has_explicit_flow_shift = params->sample.flow_shift > 0.0f && + std::isfinite(params->sample.flow_shift); float flow_shift = params->sample.flow_shift; - if (!(flow_shift > 0.0f) || !std::isfinite(flow_shift)) { - flow_shift = flux_runner_->flux_params.guidance_embed ? 1.15f : 1.0f; - } const float distilled_guidance = params->sample.distilled_guidance != 0.0f ? params->sample.distilled_guidance : 3.5f; @@ -899,9 +1017,22 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params } rng->manual_seed(static_cast(seed + batch_index)); - sd::Tensor init_latent = sd::zeros({latent_w, latent_h, 16, 1}); + const int latent_channels = flux_pipeline_latent_channels(version_); + sd::Tensor init_latent = sd::zeros({latent_w, latent_h, latent_channels, 1}); sd::Tensor noise = sd::Tensor::randn(init_latent.shape(), rng); - std::vector sigmas = ed_flux_discrete_sigmas(steps, flow_shift); + const int image_seq_len = (latent_w / patch_size) * (latent_h / patch_size); + float flux2_mu = 0.0f; + const bool use_flux2_scheduler = ed_version_is_flux2(version_) && !has_explicit_flow_shift; + std::vector sigmas; + if (use_flux2_scheduler) { + sigmas = ed_flux2_sigmas(steps, image_seq_len, &flux2_mu); + flow_shift = flux2_mu; + } else { + if (!has_explicit_flow_shift) { + flow_shift = flux_runner_->flux_params.guidance_embed ? 1.15f : 1.0f; + } + sigmas = ed_flux_discrete_sigmas(steps, flow_shift); + } if (sigmas.size() < 2) { if (error != nullptr) { *error = "failed to create Flux sigma schedule"; @@ -909,16 +1040,30 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params return false; } - LOG_INFO("flux txt2img: %dx%d latent=%dx%d steps=%d shift=%.2f guidance=%.2f cfg=%.2f seed=%" PRId64, - params->width, - params->height, - latent_w, - latent_h, - steps, - flow_shift, - distilled_guidance, - cfg_scale, - seed + batch_index); + if (use_flux2_scheduler) { + LOG_INFO("flux txt2img: %dx%d latent=%dx%d image_seq_len=%d steps=%d flux2_mu=%.3f guidance=%.2f cfg=%.2f seed=%" PRId64, + params->width, + params->height, + latent_w, + latent_h, + image_seq_len, + steps, + flux2_mu, + distilled_guidance, + cfg_scale, + seed + batch_index); + } else { + LOG_INFO("flux txt2img: %dx%d latent=%dx%d steps=%d shift=%.2f guidance=%.2f cfg=%.2f seed=%" PRId64, + params->width, + params->height, + latent_w, + latent_h, + steps, + flow_shift, + distilled_guidance, + cfg_scale, + seed + batch_index); + } sd::Tensor x = init_latent * (1.0f - sigmas[0]) + noise * sigmas[0]; sd::Tensor denoised = x; diff --git a/third_party/ggml b/third_party/ggml index e621959..019f9bb 160000 --- a/third_party/ggml +++ b/third_party/ggml @@ -1 +1 @@ -Subproject commit e621959f6baea1670d407f801c62a8cde253c9a7 +Subproject commit 019f9bbfa10054def9e3b5d84f55b39f9ce0caa5 diff --git a/tools/bench_diffusers_flux_transformer_ulysses.py b/tools/bench_diffusers_flux_transformer_ulysses.py new file mode 100644 index 0000000..0e95558 --- /dev/null +++ b/tools/bench_diffusers_flux_transformer_ulysses.py @@ -0,0 +1,230 @@ +import argparse +import json +import os +import time + +import torch +import torch.distributed as dist +import torch.profiler +from diffusers.models._modeling_parallel import ContextParallelConfig +from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel + + +def sync(): + torch.cuda.synchronize() + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + +def make_flux_ids(img_seq, txt_seq, device): + side = int(img_seq**0.5) + if side * side != img_seq: + raise ValueError(f"img_seq must be square for this benchmark, got {img_seq}") + img_ids = torch.zeros((img_seq, 3), device=device, dtype=torch.float32) + rows = torch.arange(side, device=device, dtype=torch.float32).repeat_interleave(side) + cols = torch.arange(side, device=device, dtype=torch.float32).repeat(side) + img_ids[:, 1] = rows + img_ids[:, 2] = cols + txt_ids = torch.zeros((txt_seq, 3), device=device, dtype=torch.float32) + return img_ids, txt_ids + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--mode", choices=["single", "ulysses"], default="single") + parser.add_argument("--steps", type=int, default=6) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--img-seq", type=int, default=4096) + parser.add_argument("--txt-seq", type=int, default=256) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16") + parser.add_argument("--attention-backend", default=None) + parser.add_argument("--profile", action="store_true") + parser.add_argument("--profile-steps", type=int, default=2) + parser.add_argument("--profile-row-limit", type=int, default=30) + parser.add_argument("--profile-sort", default="cuda_time_total") + args = parser.parse_args() + + distributed = args.mode == "ulysses" + if distributed: + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + else: + rank = 0 + world = 1 + local_rank = 0 + + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + dtype = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + }[args.dtype] + + t0 = time.perf_counter() + model = FluxTransformer2DModel.from_pretrained( + args.model, + subfolder="transformer", + torch_dtype=dtype, + ).to(device) + model.eval() + if args.attention_backend: + model.set_attention_backend(args.attention_backend) + if distributed: + model.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world)) + t1 = time.perf_counter() + + g = torch.Generator(device=device).manual_seed(args.seed + rank) + hidden_states = torch.randn((1, args.img_seq, model.config.in_channels), device=device, dtype=dtype, generator=g) + encoder_hidden_states = torch.randn( + (1, args.txt_seq, model.config.joint_attention_dim), + device=device, + dtype=dtype, + generator=g, + ) + pooled = torch.randn((1, model.config.pooled_projection_dim), device=device, dtype=dtype, generator=g) + timestep = torch.ones((1,), device=device, dtype=dtype) + guidance = torch.full((1,), 3.5, device=device, dtype=dtype) if model.config.guidance_embeds else None + img_ids, txt_ids = make_flux_ids(args.img_seq, args.txt_seq, device) + + def run_once(): + with torch.inference_mode(): + out = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + guidance=guidance, + return_dict=False, + )[0] + return out + + for _ in range(args.warmup): + _ = run_once() + sync() + + times = [] + total_start = time.perf_counter() + for _ in range(args.steps): + sync() + start = time.perf_counter() + out = run_once() + sync() + end = time.perf_counter() + times.append((end - start) * 1000.0) + total_end = time.perf_counter() + + checksum = float(out.float().mean().detach().cpu()) + payload = { + "mode": args.mode, + "rank": rank, + "world": world, + "local_rank": local_rank, + "dtype": args.dtype, + "attention_backend": args.attention_backend, + "load_s": t1 - t0, + "steps": args.steps, + "warmup": args.warmup, + "img_seq": args.img_seq, + "txt_seq": args.txt_seq, + "times_ms": times, + "mean_ms": sum(times) / len(times), + "total_s": total_end - total_start, + "checksum": checksum, + } + print("DIFFUSERS_TRANSFORMER_PROFILE " + json.dumps(payload, sort_keys=True), flush=True) + + if args.profile: + sync() + activities = [torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + with torch.profiler.profile( + activities=activities, + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + for step in range(args.profile_steps): + sync() + with torch.profiler.record_function(f"diffusers_transformer_forward_{step}"): + _ = run_once() + sync() + prof.step() + sync() + + events = list(prof.key_averages()) + + def event_us(event, *names): + for name in names: + value = getattr(event, name, None) + if value is not None: + return float(value) + return 0.0 + + def event_count(event): + return int(getattr(event, "count", 0)) + + def event_payload(event): + return { + "name": event.key, + "count": event_count(event), + "cpu_total_ms": event_us(event, "cpu_time_total") / 1000.0, + "cpu_self_ms": event_us(event, "self_cpu_time_total") / 1000.0, + "cuda_total_ms": event_us(event, "cuda_time_total", "device_time_total") / 1000.0, + "cuda_self_ms": event_us(event, "self_cuda_time_total", "self_device_time_total") / 1000.0, + } + + top_events = sorted( + events, + key=lambda event: event_us(event, args.profile_sort, "cuda_time_total", "device_time_total"), + reverse=True, + )[: args.profile_row_limit] + keywords = ( + "nccl", + "all_to_all", + "alltoall", + "all_gather", + "allgather", + "send", + "recv", + "scaled_dot_product", + "flash", + "attention", + ) + selected_events = [ + event + for event in events + if any(keyword in event.key.lower() for keyword in keywords) + ] + selected_events = sorted( + selected_events, + key=lambda event: event_us(event, "cuda_time_total", "device_time_total", "cpu_time_total"), + reverse=True, + ) + profile_payload = { + "mode": args.mode, + "rank": rank, + "world": world, + "steps": args.profile_steps, + "sort": args.profile_sort, + "top": [event_payload(event) for event in top_events], + "selected": [event_payload(event) for event in selected_events[: args.profile_row_limit]], + } + print("DIFFUSERS_TRANSFORMER_TORCH_PROFILE " + json.dumps(profile_payload, sort_keys=True), flush=True) + print( + f"DIFFUSERS_TRANSFORMER_TORCH_PROFILE_TABLE rank={rank} world={world} sort={args.profile_sort}", + flush=True, + ) + print(prof.key_averages().table(sort_by=args.profile_sort, row_limit=args.profile_row_limit), flush=True) + + if distributed: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/bench_diffusers_flux_ulysses.py b/tools/bench_diffusers_flux_ulysses.py new file mode 100644 index 0000000..da387e1 --- /dev/null +++ b/tools/bench_diffusers_flux_ulysses.py @@ -0,0 +1,120 @@ +import argparse +import json +import os +import time +from pathlib import Path + +import torch +import torch.distributed as dist +from diffusers import FluxPipeline +from diffusers.models._modeling_parallel import ContextParallelConfig + + +def synchronize(): + if torch.cuda.is_available(): + torch.cuda.synchronize() + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--steps", type=int, default=6) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--guidance", type=float, default=3.5) + parser.add_argument("--mode", choices=["single", "ulysses"], default="single") + parser.add_argument("--output", default="/tmp/diffusers_flux.png") + parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16") + parser.add_argument("--max-sequence-length", type=int, default=256) + args = parser.parse_args() + + distributed = args.mode == "ulysses" + if distributed: + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + else: + rank = 0 + world = 1 + local_rank = 0 + + torch.cuda.set_device(local_rank) + dtype = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + }[args.dtype] + + t0 = time.perf_counter() + pipe = FluxPipeline.from_pretrained(args.model, torch_dtype=dtype) + pipe.to(f"cuda:{local_rank}") + pipe.set_progress_bar_config(disable=True) + t1 = time.perf_counter() + + if distributed: + cp_config = ContextParallelConfig(ring_degree=1, ulysses_degree=world) + pipe.transformer.enable_parallelism(config=cp_config) + + transformer_times = [] + orig_forward = pipe.transformer.forward + + def timed_forward(*f_args, **f_kwargs): + synchronize() + s = time.perf_counter() + out = orig_forward(*f_args, **f_kwargs) + synchronize() + e = time.perf_counter() + transformer_times.append((e - s) * 1000.0) + return out + + pipe.transformer.forward = timed_forward + + generator = torch.Generator(device=f"cuda:{local_rank}").manual_seed(args.seed) + synchronize() + infer_start = time.perf_counter() + result = pipe( + prompt=args.prompt, + width=args.width, + height=args.height, + num_inference_steps=args.steps, + guidance_scale=args.guidance, + generator=generator, + output_type="pil" if rank == 0 else "latent", + max_sequence_length=args.max_sequence_length, + ) + synchronize() + infer_end = time.perf_counter() + + if rank == 0 and hasattr(result, "images") and result.images is not None and args.output: + image = result.images[0] + if hasattr(image, "save"): + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + image.save(args.output) + + payload = { + "mode": args.mode, + "rank": rank, + "world": world, + "local_rank": local_rank, + "dtype": args.dtype, + "width": args.width, + "height": args.height, + "steps": args.steps, + "load_s": t1 - t0, + "inference_s": infer_end - infer_start, + "transformer_ms": transformer_times, + "transformer_sum_ms": sum(transformer_times), + } + print("DIFFUSERS_PROFILE " + json.dumps(payload, sort_keys=True), flush=True) + + if distributed: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/bench_diffusers_wan_transformer.py b/tools/bench_diffusers_wan_transformer.py new file mode 100644 index 0000000..5ae7e02 --- /dev/null +++ b/tools/bench_diffusers_wan_transformer.py @@ -0,0 +1,233 @@ +import argparse +import json +import os +import time + +import torch +import torch.distributed as dist +import torch.profiler +from diffusers import WanTransformer3DModel +from diffusers.models._modeling_parallel import ContextParallelConfig + + +def sync(): + torch.cuda.synchronize() + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--mode", choices=["single", "ulysses"], default="single") + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--latent-frames", type=int, default=10) + parser.add_argument("--steps", type=int, default=4) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--dtype", choices=["fp16", "bf16", "fp32"], default="fp16") + parser.add_argument("--attention-backend", default=None) + parser.add_argument("--fuse-qkv", action="store_true") + parser.add_argument("--text-seq", type=int, default=512) + parser.add_argument("--profile", action="store_true") + parser.add_argument("--profile-steps", type=int, default=1) + parser.add_argument("--profile-row-limit", type=int, default=30) + parser.add_argument("--profile-sort", default="cuda_time_total") + args = parser.parse_args() + + distributed = args.mode == "ulysses" + if distributed: + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + else: + rank = 0 + world = 1 + local_rank = 0 + + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + dtype = { + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32, + }[args.dtype] + + t0 = time.perf_counter() + model = WanTransformer3DModel.from_pretrained( + args.model, + subfolder="transformer", + local_files_only=True, + low_cpu_mem_usage=True, + ) + model.to(device=device, dtype=dtype) + model.eval() + if args.attention_backend: + model.set_attention_backend(args.attention_backend) + if args.fuse_qkv: + model.fuse_qkv_projections() + if distributed: + model.enable_parallelism(config=ContextParallelConfig(ring_degree=1, ulysses_degree=world)) + t1 = time.perf_counter() + + latent_h = args.height // 8 + latent_w = args.width // 8 + g = torch.Generator(device=device).manual_seed(args.seed) + hidden_states = torch.randn( + (1, model.config.in_channels, args.latent_frames, latent_h, latent_w), + device=device, + dtype=dtype, + generator=g, + ) + prompt_embeds = torch.randn( + (1, args.text_seq, model.config.text_dim), + device=device, + dtype=dtype, + generator=g, + ) + negative_prompt_embeds = torch.randn( + (1, args.text_seq, model.config.text_dim), + device=device, + dtype=dtype, + generator=g, + ) + timestep = torch.full((1,), 999.0, device=device, dtype=torch.float32) + + def forward_pair(): + with torch.inference_mode(): + cond = model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + return_dict=False, + )[0] + uncond = model( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + return_dict=False, + )[0] + out = uncond + 5.0 * (cond - uncond) + return out + + for _ in range(args.warmup): + _ = forward_pair() + sync() + + times = [] + total_start = time.perf_counter() + for _ in range(args.steps): + sync() + start = time.perf_counter() + out = forward_pair() + sync() + end = time.perf_counter() + times.append((end - start) * 1000.0) + total_end = time.perf_counter() + + payload = { + "width": args.width, + "height": args.height, + "latent_frames": args.latent_frames, + "latent_h": latent_h, + "latent_w": latent_w, + "seq": (args.latent_frames // model.config.patch_size[0]) + * (latent_h // model.config.patch_size[1]) + * (latent_w // model.config.patch_size[2]), + "dtype": args.dtype, + "param_dtype": str(next(model.parameters()).dtype), + "attention_backend": args.attention_backend, + "fuse_qkv": args.fuse_qkv, + "load_s": t1 - t0, + "steps": args.steps, + "warmup": args.warmup, + "pair_times_ms": times, + "mean_pair_ms": sum(times) / len(times), + "total_s": total_end - total_start, + "checksum": float(out.float().mean().detach().cpu()), + } + print("DIFFUSERS_WAN_TRANSFORMER_PROFILE " + json.dumps(payload, sort_keys=True), flush=True) + + if args.profile: + activities = [torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + with torch.profiler.profile( + activities=activities, + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + for _ in range(args.profile_steps): + sync() + _ = forward_pair() + sync() + prof.step() + events = list(prof.key_averages()) + + def event_us(event, *names): + for name in names: + value = getattr(event, name, None) + if value is not None: + return float(value) + return 0.0 + + def event_payload(event): + return { + "name": event.key, + "count": int(getattr(event, "count", 0)), + "cpu_total_ms": event_us(event, "cpu_time_total") / 1000.0, + "cpu_self_ms": event_us(event, "self_cpu_time_total") / 1000.0, + "cuda_total_ms": event_us(event, "cuda_time_total", "device_time_total") / 1000.0, + "cuda_self_ms": event_us(event, "self_cuda_time_total", "self_device_time_total") / 1000.0, + } + + top_events = sorted( + events, + key=lambda event: event_us(event, args.profile_sort, "cuda_time_total", "device_time_total"), + reverse=True, + )[: args.profile_row_limit] + keywords = ( + "nccl", + "all_to_all", + "alltoall", + "all_gather", + "allgather", + "send", + "recv", + "scaled_dot_product", + "flash", + "attention", + ) + selected_events = [ + event + for event in events + if any(keyword in event.key.lower() for keyword in keywords) + ] + selected_events = sorted( + selected_events, + key=lambda event: event_us(event, "cuda_time_total", "device_time_total", "cpu_time_total"), + reverse=True, + ) + profile_payload = { + "mode": args.mode, + "rank": rank, + "world": world, + "steps": args.profile_steps, + "sort": args.profile_sort, + "top": [event_payload(event) for event in top_events], + "selected": [event_payload(event) for event in selected_events[: args.profile_row_limit]], + } + print("DIFFUSERS_WAN_TRANSFORMER_TORCH_PROFILE " + json.dumps(profile_payload, sort_keys=True), flush=True) + print( + f"DIFFUSERS_WAN_TRANSFORMER_TORCH_PROFILE_TABLE rank={rank} world={world} sort={args.profile_sort}", + flush=True, + ) + print(prof.key_averages().table(sort_by=args.profile_sort, row_limit=args.profile_row_limit), flush=True) + + if distributed: + dist.destroy_process_group() + + +if __name__ == "__main__": + main()