From cf5be40a321fe025069fde2b74aac94c666e324b Mon Sep 17 00:00:00 2001 From: WenQ7 <136597310+WenQ7@users.noreply.github.com> Date: Mon, 6 Jul 2026 02:26:48 +0000 Subject: [PATCH] feat: add DCU support for LongCat and Flux pipelines via C API attention --- xllm/models/dit/autoencoders/autoencoder_kl.h | 9 ++ xllm/models/dit/pipelines/pipeline_flux.h | 1 + .../dit/pipelines/pipeline_flux_control.h | 6 + .../models/dit/pipelines/pipeline_flux_fill.h | 17 ++- .../dit/pipelines/pipeline_longcat_audiodit.h | 24 +++- .../dit/pipelines/pipeline_longcat_image.h | 45 ++++-- .../pipelines/pipeline_longcat_image_edit.h | 74 ++++++---- .../dit/transformers/transformer_flux.h | 86 +++++++++++- .../transformer_longcat_audiodit.h | 129 ++++++++++++++++++ xllm/models/models.h | 4 + 10 files changed, 343 insertions(+), 52 deletions(-) diff --git a/xllm/models/dit/autoencoders/autoencoder_kl.h b/xllm/models/dit/autoencoders/autoencoder_kl.h index a237491962..847979f20c 100644 --- a/xllm/models/dit/autoencoders/autoencoder_kl.h +++ b/xllm/models/dit/autoencoders/autoencoder_kl.h @@ -1024,6 +1024,15 @@ class VAEImpl : public torch::nn::Module { return posterior.sample(seed); } + torch::Tensor encode_mode(const torch::Tensor& images) { + auto enc = encoder_(images); + if (args_.use_quant_conv()) { + enc = quant_conv_(enc); + } + auto posterior = DiagonalGaussianDistribution(enc); + return posterior.mode(); + } + torch::Tensor decode(const torch::Tensor& latents) { torch::Tensor processed_latents = latents; diff --git a/xllm/models/dit/pipelines/pipeline_flux.h b/xllm/models/dit/pipelines/pipeline_flux.h index 60d53d8e45..c0c517b099 100644 --- a/xllm/models/dit/pipelines/pipeline_flux.h +++ b/xllm/models/dit/pipelines/pipeline_flux.h @@ -336,4 +336,5 @@ class FluxPipelineImpl : public FluxPipelineBaseImpl { TORCH_MODULE(FluxPipeline); REGISTER_DIT_MODEL(flux, FluxPipeline); +REGISTER_DIT_MODEL_WITH_VARNAME(flux_pipeline, FluxPipeline, FluxPipeline); } // namespace xllm diff --git a/xllm/models/dit/pipelines/pipeline_flux_control.h b/xllm/models/dit/pipelines/pipeline_flux_control.h index 4710648516..80eaebe57f 100644 --- a/xllm/models/dit/pipelines/pipeline_flux_control.h +++ b/xllm/models/dit/pipelines/pipeline_flux_control.h @@ -81,6 +81,9 @@ class FluxControlPipelineImpl : public FluxPipelineBaseImpl { ? std::make_optional(input.pooled_prompt_embeds) : std::nullopt; + CHECK(control_image.defined()) + << "FluxControl pipeline requires an input control_image."; + auto output = forward_impl(prompts, prompts_2, control_image, @@ -351,4 +354,7 @@ class FluxControlPipelineImpl : public FluxPipelineBaseImpl { TORCH_MODULE(FluxControlPipeline); REGISTER_DIT_MODEL(flux_control, FluxControlPipeline); +REGISTER_DIT_MODEL_WITH_VARNAME(flux_control_pipeline, + FluxControlPipeline, + FluxControlPipeline); } // namespace xllm diff --git a/xllm/models/dit/pipelines/pipeline_flux_fill.h b/xllm/models/dit/pipelines/pipeline_flux_fill.h index e58c471288..ee2d1b9993 100644 --- a/xllm/models/dit/pipelines/pipeline_flux_fill.h +++ b/xllm/models/dit/pipelines/pipeline_flux_fill.h @@ -78,8 +78,12 @@ class FluxFillPipelineImpl : public FluxPipelineBaseImpl { ? std::nullopt : std::make_optional(input.prompts_2); - auto image = input.images.defined() ? std::make_optional(input.images) - : std::nullopt; + std::optional image = std::nullopt; + if (input.images.defined()) { + image = input.images; + } else if (!input.images_list.empty()) { + image = input.images_list[0]; + } auto mask_image = input.mask_images.defined() ? std::make_optional(input.mask_images) : std::nullopt; @@ -98,6 +102,12 @@ class FluxFillPipelineImpl : public FluxPipelineBaseImpl { ? std::make_optional(input.pooled_prompt_embeds) : std::nullopt; + CHECK(image.has_value()) + << "FluxFill pipeline requires an input image in images or " + "images_list."; + CHECK(mask_image.has_value()) + << "FluxFill pipeline requires an input mask_image."; + auto output = forward_impl(prompts, prompts_2, image, @@ -424,4 +434,7 @@ class FluxFillPipelineImpl : public FluxPipelineBaseImpl { TORCH_MODULE(FluxFillPipeline); REGISTER_DIT_MODEL(fluxfill, FluxFillPipeline); +REGISTER_DIT_MODEL_WITH_VARNAME(flux_fill_pipeline, + FluxFillPipeline, + FluxFillPipeline); } // namespace xllm diff --git a/xllm/models/dit/pipelines/pipeline_longcat_audiodit.h b/xllm/models/dit/pipelines/pipeline_longcat_audiodit.h index 48b25a7b03..dac007839d 100644 --- a/xllm/models/dit/pipelines/pipeline_longcat_audiodit.h +++ b/xllm/models/dit/pipelines/pipeline_longcat_audiodit.h @@ -500,7 +500,11 @@ class LongCatAudioDiTPipelineImpl final : public torch::nn::Module { load_module_from_state_dicts(*vae_loader, vae_.ptr().get()); vae_->to(options_.device()); +#if defined(USE_DCU) + vae_->to(torch::kFloat32); +#else vae_->to_half(); +#endif text_encoder_->load_model(std::move(text_encoder_loader)); text_encoder_->to(options_.device()); @@ -521,7 +525,11 @@ class LongCatAudioDiTPipelineImpl final : public torch::nn::Module { load_module_from_state_dicts(*flat_loader, vae_.ptr().get(), "vae."); vae_->to(options_.device()); +#if defined(USE_DCU) + vae_->to(torch::kFloat32); +#else vae_->to_half(); +#endif text_encoder_->load_model_from_state_dicts( flat_loader->get_state_dicts()); @@ -605,13 +613,15 @@ TORCH_MODULE(LongCatAudioDiTPipeline); // Register under "audiodit" to match the model_type field in config.json. namespace { const bool longcat_audio_dit_registered = []() { - ModelRegistry::register_dit_model_factory( - "audiodit", [](const DiTModelContext& context) { - LongCatAudioDiTPipeline model(context); - model->eval(); - return std::make_unique>( - std::move(model), context.get_tensor_options()); - }); + auto factory = [](const DiTModelContext& context) { + LongCatAudioDiTPipeline model(context); + model->eval(); + return std::make_unique>( + std::move(model), context.get_tensor_options()); + }; + ModelRegistry::register_dit_model_factory("audiodit", factory); + ModelRegistry::register_dit_model_factory("LongCat-AudioDiT", factory); + ModelRegistry::register_dit_model_factory("LongCatAudioDiTPipeline", factory); return true; }(); } // namespace diff --git a/xllm/models/dit/pipelines/pipeline_longcat_image.h b/xllm/models/dit/pipelines/pipeline_longcat_image.h index ced68a8d30..d656b03258 100644 --- a/xllm/models/dit/pipelines/pipeline_longcat_image.h +++ b/xllm/models/dit/pipelines/pipeline_longcat_image.h @@ -55,6 +55,13 @@ constexpr const char* PROMPT_TEMPLATE_ENCODE_PREFIX = constexpr const char* PROMPT_TEMPLATE_ENCODE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"; +inline int32_t get_longcat_vlm_tp_port(int32_t base_port, int32_t rank) { + const int32_t port = base_port + rank; + CHECK_LE(port, 65535) << "LongCat VLM ProcessGroup port overflow: base_port=" + << base_port << ", rank=" << rank; + return port; +} + inline float calculate_longcat_shift(int64_t image_seq_len, int64_t base_seq_len = 256, int64_t max_seq_len = 4096, @@ -134,15 +141,19 @@ class LongCatImagePipelineImpl : public torch::nn::Module { const auto& original_parallel_args = context.get_parallel_args(); ParallelArgs vlm_parallel_args = original_parallel_args; if (original_parallel_args.tp_group_ == nullptr) { + constexpr int32_t kLongCatImageVlmTpBasePort = 29500; + const int32_t vlm_tp_port = get_longcat_vlm_tp_port( + kLongCatImageVlmTpBasePort, original_parallel_args.rank()); LOG(INFO) - << "Creating real ProcessGroup for single-device VLM initialization."; - vlm_tp_group_ = create_process_group(0, - 1, - 1, - 29500, - false, - "127.0.0.1", - "vlm_tp_group", + << "Creating real ProcessGroup for single-device VLM initialization " + << "on 127.0.0.1:" << vlm_tp_port; + vlm_tp_group_ = create_process_group(/*rank=*/0, + /*world_size=*/1, + /*rank_size=*/1, + vlm_tp_port, + /*trans=*/false, + /*host=*/"127.0.0.1", + /*group_name=*/"vlm_tp_group", options_.device()); vlm_parallel_args.tp_group_ = vlm_tp_group_.get(); } @@ -637,6 +648,9 @@ class LongCatImagePipelineImpl : public torch::nn::Module { const torch::Tensor& positions, const torch::Tensor& attention_mask) { ModelInputParams params; +#if defined(USE_DCU) + params.graph.use_dense_flash_attention = true; +#endif int64_t actual_seq_len; if (positions.dim() == 2) { @@ -866,13 +880,14 @@ TORCH_MODULE(LongCatImagePipeline); // code. namespace { const bool longcat_image_dit_registered = []() { - ModelRegistry::register_dit_model_factory( - "LongCat-Image", [](const DiTModelContext& context) { - LongCatImagePipeline model(context); - model->eval(); - return std::make_unique>( - std::move(model), context.get_tensor_options()); - }); + auto factory = [](const DiTModelContext& context) { + LongCatImagePipeline model(context); + model->eval(); + return std::make_unique>( + std::move(model), context.get_tensor_options()); + }; + ModelRegistry::register_dit_model_factory("LongCat-Image", factory); + ModelRegistry::register_dit_model_factory("LongCatImagePipeline", factory); return true; }(); } // namespace diff --git a/xllm/models/dit/pipelines/pipeline_longcat_image_edit.h b/xllm/models/dit/pipelines/pipeline_longcat_image_edit.h index c19b8272b2..9bbb4fb2db 100644 --- a/xllm/models/dit/pipelines/pipeline_longcat_image_edit.h +++ b/xllm/models/dit/pipelines/pipeline_longcat_image_edit.h @@ -32,7 +32,9 @@ #include "core/framework/request/dit_request_state.h" #include "core/framework/tokenizer/tokenizer.h" #include "core/layers/common/attention_metadata_builder.h" +#if defined(USE_CUDA) #include "core/layers/cuda/flashinfer_workspace.h" +#endif #include "models/dit/autoencoders/autoencoder_kl.h" #include "models/dit/pipelines/pipeline_longcat_image.h" #include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h" @@ -80,7 +82,7 @@ inline std::pair calculate_dimensions_edit( return {w, h}; } -// LongCat-Image-Edit pipeline for CUDA backend. +// LongCat-Image-Edit pipeline. class LongCatImageEditPipelineImpl : public torch::nn::Module { public: explicit LongCatImageEditPipelineImpl(const DiTModelContext& context) @@ -89,9 +91,11 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module { const auto& model_args = context.get_model_args("vae"); options_ = context.get_tensor_options(); - // Initialize FlashinferWorkspace for attention operations +#if defined(USE_CUDA) + // Initialize FlashinferWorkspace for CUDA attention operations. layer::flashinfer::FlashinferWorkspace::get_instance().initialize( options_.device()); +#endif vae_scale_factor_ = 1 << (model_args.block_out_channels().size() - 1); vae_shift_factor_ = model_args.shift_factor(); @@ -129,16 +133,21 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module { const auto& original_parallel_args = context.get_parallel_args(); ParallelArgs vlm_parallel_args = original_parallel_args; if (original_parallel_args.tp_group_ == nullptr) { + constexpr int32_t kLongCatImageEditVlmTpBasePort = 29564; + const int32_t vlm_tp_port = get_longcat_vlm_tp_port( + kLongCatImageEditVlmTpBasePort, original_parallel_args.rank()); LOG(INFO) - << "Creating real ProcessGroup for single-device VLM initialization."; - vlm_tp_group_ = create_process_group(0, - 1, - 1, - 29501, // Different port from T2I - false, - "127.0.0.1", - "vlm_tp_group_longcat_edit", - options_.device()); + << "Creating real ProcessGroup for single-device VLM initialization " + << "on 127.0.0.1:" << vlm_tp_port; + vlm_tp_group_ = + create_process_group(/*rank=*/0, + /*world_size=*/1, + /*rank_size=*/1, + vlm_tp_port, + /*trans=*/false, + /*host=*/"127.0.0.1", + /*group_name=*/"vlm_tp_group_longcat_edit", + options_.device()); vlm_parallel_args.tp_group_ = vlm_tp_group_.get(); } @@ -200,10 +209,22 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module { ? std::make_optional(input.negative_pooled_prompt_embeds) : std::nullopt; - auto image = input.images.defined() ? std::make_optional(input.images) - : std::nullopt; + std::optional image = std::nullopt; + if (!input.images_list.empty()) { + CHECK_EQ(input.images_list.size(), 1) + << "LongCat-Image-Edit expects exactly one input image tensor, got " + << input.images_list.size(); + image = input.images_list[0]; + } else if (input.images.defined()) { + image = input.images; + } - CHECK(image.has_value()) << "LongCat-Image-Edit requires an input image."; + CHECK(image.has_value()) + << "LongCat-Image-Edit requires an input image in images or " + "images_list. batch_size=" + << input.batch_size << ", prompts=" << input.prompts.size() + << ", images_defined=" << input.images.defined() + << ", images_list_size=" << input.images_list.size(); std::vector output = forward_( image, // input image @@ -586,8 +607,8 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module { return {prompt_embeds.value(), text_ids}; } - torch::Tensor encode_vae_image(const torch::Tensor& image, int64_t seed) { - torch::Tensor latents = vae_->encode(image, seed); + torch::Tensor encode_vae_image(const torch::Tensor& image) { + torch::Tensor latents = vae_->encode_mode(image); latents = (latents - vae_shift_factor_) * vae_scaling_factor_; return latents; } @@ -611,7 +632,7 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module { torch::Tensor image_latents; if (image.size(1) != num_channels_latents) { - image_latents = encode_vae_image(image, seed); + image_latents = encode_vae_image(image); } else { image_latents = image; } @@ -799,6 +820,9 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module { const torch::Tensor& attention_mask, std::optional mm_data_opt = std::nullopt) { ModelInputParams params; +#if defined(USE_DCU) + params.graph.use_dense_flash_attention = true; +#endif int64_t actual_seq_len; if (positions.dim() == 2) { @@ -1123,13 +1147,15 @@ TORCH_MODULE(LongCatImageEditPipeline); // Register LongCat-Image-Edit as DiT model. namespace { const bool longcat_image_edit_dit_registered = []() { - ModelRegistry::register_dit_model_factory( - "LongCat-Image-Edit", [](const DiTModelContext& context) { - LongCatImageEditPipeline model(context); - model->eval(); - return std::make_unique>( - std::move(model), context.get_tensor_options()); - }); + auto factory = [](const DiTModelContext& context) { + LongCatImageEditPipeline model(context); + model->eval(); + return std::make_unique>( + std::move(model), context.get_tensor_options()); + }; + ModelRegistry::register_dit_model_factory("LongCat-Image-Edit", factory); + ModelRegistry::register_dit_model_factory("LongCatImageEditPipeline", + factory); return true; }(); } // namespace diff --git a/xllm/models/dit/transformers/transformer_flux.h b/xllm/models/dit/transformers/transformer_flux.h index 9a16c96835..5b1038b38b 100644 --- a/xllm/models/dit/transformers/transformer_flux.h +++ b/xllm/models/dit/transformers/transformer_flux.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -33,6 +34,9 @@ limitations under the License. #include "core/framework/state_dict/utils.h" #include "core/layers/common/add_matmul.h" #include "core/layers/common/rms_norm.h" +#if defined(USE_DCU) +#include "core/layers/dcu/flash_attention.h" +#endif #include "framework/model_context.h" #include "models/model_registry.h" #if defined(USE_NPU) @@ -80,6 +84,80 @@ inline torch::Tensor apply_rotary_emb(const torch::Tensor& x, #endif } +inline torch::Tensor flux_scaled_dot_product_attention( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value) { +#if defined(USE_DCU) + static thread_local bool logged = false; + if (!logged) { + LOG(INFO) << "Flux attention uses dense varlen flash attention on DCU."; + logged = true; + } + + CHECK_EQ(query.dim(), 4) << "Flux query must be [B,H,S,D]"; + CHECK_EQ(key.dim(), 4) << "Flux key must be [B,H,S,D]"; + CHECK_EQ(value.dim(), 4) << "Flux value must be [B,H,S,D]"; + CHECK_EQ(query.size(0), key.size(0)) << "Flux q/k batch mismatch"; + CHECK_EQ(query.size(0), value.size(0)) << "Flux q/v batch mismatch"; + CHECK_EQ(query.size(1), key.size(1)) << "Flux q/k head mismatch"; + CHECK_EQ(key.size(1), value.size(1)) << "Flux k/v head mismatch"; + CHECK_EQ(query.size(3), key.size(3)) << "Flux q/k head dim mismatch"; + CHECK_EQ(query.size(3), value.size(3)) << "Flux q/v head dim mismatch"; + CHECK_EQ(key.size(2), value.size(2)) << "Flux k/v seq mismatch"; + + const int64_t batch_size = query.size(0); + const int64_t num_heads = query.size(1); + const int64_t q_seq_len = query.size(2); + const int64_t kv_seq_len = key.size(2); + const int64_t head_dim = query.size(3); + + const auto output_dtype = query.scalar_type(); + torch::Tensor attn_query = query; + torch::Tensor attn_key = key; + torch::Tensor attn_value = value; + if (output_dtype == at::kFloat) { + attn_query = query.to(torch::kBFloat16); + attn_key = key.to(torch::kBFloat16); + attn_value = value.to(torch::kBFloat16); + } + + auto cu_options = + torch::TensorOptions().dtype(torch::kInt32).device(query.device()); + auto cu_seqlens_q = + torch::arange(batch_size + 1, cu_options).mul_(q_seq_len).contiguous(); + auto cu_seqlens_k = + torch::arange(batch_size + 1, cu_options).mul_(kv_seq_len).contiguous(); + + auto dense_query = attn_query.transpose(1, 2).contiguous().view( + {batch_size * q_seq_len, num_heads, head_dim}); + auto dense_key = attn_key.transpose(1, 2).contiguous().view( + {batch_size * kv_seq_len, key.size(1), head_dim}); + auto dense_value = attn_value.transpose(1, 2).contiguous().view( + {batch_size * kv_seq_len, value.size(1), head_dim}); + + auto output = layer::dense_varlen_flash_attention( + dense_query, + dense_key, + dense_value, + cu_seqlens_q, + cu_seqlens_k, + std::pow(static_cast(head_dim), -0.5), + /*is_causal=*/false); + return output.view({batch_size, q_seq_len, num_heads, head_dim}) + .transpose(1, 2) + .contiguous() + .to(output_dtype); +#endif + + return torch::scaled_dot_product_attention(query, + key, + value, + torch::nullopt, + /*dropout_p=*/0.0, + /*is_causal=*/false); +} + // Helper function for rotary position embedding torch::Tensor get_1d_rotary_pos_embed( int64_t dim, @@ -290,8 +368,8 @@ class FluxSingleAttentionImpl : public torch::nn::Module { query = apply_rotary_emb(query, image_rotary_emb); key = apply_rotary_emb(key, image_rotary_emb); // Compute scaled dot-product attention (no mask, no dropout) - torch::Tensor attn_output = torch::scaled_dot_product_attention( - query, key, value, torch::nullopt, 0.0, false); + torch::Tensor attn_output = + flux_scaled_dot_product_attention(query, key, value); attn_output = attn_output.to(query.dtype()); return attn_output.transpose(1, 2).flatten(2); #else @@ -465,8 +543,8 @@ class FluxAttentionImpl : public torch::nn::Module { query1 = query1.transpose(1, 2); key1 = key1.transpose(1, 2); value1 = value1.transpose(1, 2); - torch::Tensor attn_output = torch::scaled_dot_product_attention( - query1, key1, value1, torch::nullopt, 0.0, false); + torch::Tensor attn_output = + flux_scaled_dot_product_attention(query1, key1, value1); attn_output = attn_output.transpose(1, 2).reshape( {batch_size, -1, attn_heads * head_dim}); #else diff --git a/xllm/models/dit/transformers/transformer_longcat_audiodit.h b/xllm/models/dit/transformers/transformer_longcat_audiodit.h index 466cbfcf57..39f95d5541 100644 --- a/xllm/models/dit/transformers/transformer_longcat_audiodit.h +++ b/xllm/models/dit/transformers/transformer_longcat_audiodit.h @@ -44,6 +44,9 @@ limitations under the License. #include "core/framework/dit_model_loader.h" #include "core/framework/model_context.h" #include "core/framework/request/dit_request_state.h" +#if defined(USE_DCU) +#include "core/layers/dcu/flash_attention.h" +#endif #include "models/dit/autoencoders/autoencoder_kl.h" // randn_tensor #include "models/dit/encoders/umt5_encoder.h" @@ -725,6 +728,124 @@ inline torch::Tensor apply_rotary_emb(const torch::Tensor& x, .to(x.dtype()); } +#if defined(USE_DCU) +inline torch::Tensor audio_dense_varlen_flash_attention( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + std::optional key_mask = std::nullopt) { + CHECK_EQ(query.dim(), 4) << "AudioDiT query must be [B,H,S,D]"; + CHECK_EQ(key.dim(), 4) << "AudioDiT key must be [B,H,S,D]"; + CHECK_EQ(value.dim(), 4) << "AudioDiT value must be [B,H,S,D]"; + CHECK_EQ(query.size(0), key.size(0)) << "AudioDiT q/k batch mismatch"; + CHECK_EQ(query.size(0), value.size(0)) << "AudioDiT q/v batch mismatch"; + CHECK_EQ(query.size(1), key.size(1)) << "AudioDiT q/k head mismatch"; + CHECK_EQ(key.size(1), value.size(1)) << "AudioDiT k/v head mismatch"; + CHECK_EQ(query.size(3), key.size(3)) << "AudioDiT q/k head dim mismatch"; + CHECK_EQ(query.size(3), value.size(3)) << "AudioDiT q/v head dim mismatch"; + CHECK_EQ(key.size(2), value.size(2)) << "AudioDiT k/v seq mismatch"; + + const int64_t batch_size = query.size(0); + const int64_t num_heads = query.size(1); + const int64_t q_seq_len = query.size(2); + const int64_t kv_seq_len = key.size(2); + const int64_t head_dim = query.size(3); + + if (key_mask.has_value() && key_mask->all().item()) { + key_mask = std::nullopt; + } + + const auto output_dtype = query.scalar_type(); + torch::Tensor attn_query = query; + torch::Tensor attn_key = key; + torch::Tensor attn_value = value; + if (output_dtype == at::kFloat) { + attn_query = query.to(torch::kBFloat16); + attn_key = key.to(torch::kBFloat16); + attn_value = value.to(torch::kBFloat16); + } + + auto dense_query = attn_query.transpose(1, 2).contiguous().view( + {batch_size * q_seq_len, num_heads, head_dim}); + std::vector cu_seqlens_q_vec; + cu_seqlens_q_vec.reserve(batch_size + 1); + for (int64_t i = 0; i <= batch_size; ++i) { + cu_seqlens_q_vec.push_back(static_cast(i * q_seq_len)); + } + + torch::Tensor dense_key; + torch::Tensor dense_value; + std::vector cu_seqlens_k_vec; + cu_seqlens_k_vec.reserve(batch_size + 1); + cu_seqlens_k_vec.push_back(0); + + if (key_mask.has_value()) { + CHECK_EQ(key_mask->dim(), 2) << "AudioDiT key mask must be [B,S]"; + CHECK_EQ(key_mask->size(0), batch_size) + << "AudioDiT key mask batch mismatch"; + CHECK_EQ(key_mask->size(1), kv_seq_len) << "AudioDiT key mask seq mismatch"; + + std::vector key_chunks; + std::vector value_chunks; + key_chunks.reserve(batch_size); + value_chunks.reserve(batch_size); + int32_t total_kv_len = 0; + for (int64_t i = 0; i < batch_size; ++i) { + auto valid_indices = + torch::nonzero(key_mask->select(0, i).to(torch::kBool)).flatten(); + const int64_t valid_len = valid_indices.numel(); + CHECK_GT(valid_len, 0) + << "AudioDiT dense flash attention requires at least one valid key"; + CHECK_LE(valid_len, + static_cast(std::numeric_limits::max()) - + total_kv_len) + << "AudioDiT key sequence length exceeds int32 range"; + total_kv_len += static_cast(valid_len); + cu_seqlens_k_vec.push_back(total_kv_len); + + key_chunks.push_back( + attn_key.select(0, i).transpose(0, 1).contiguous().index_select( + 0, valid_indices)); + value_chunks.push_back( + attn_value.select(0, i).transpose(0, 1).contiguous().index_select( + 0, valid_indices)); + } + dense_key = torch::cat(key_chunks, /*dim=*/0); + dense_value = torch::cat(value_chunks, /*dim=*/0); + } else { + dense_key = attn_key.transpose(1, 2).contiguous().view( + {batch_size * kv_seq_len, attn_key.size(1), head_dim}); + dense_value = attn_value.transpose(1, 2).contiguous().view( + {batch_size * kv_seq_len, attn_value.size(1), head_dim}); + for (int64_t i = 1; i <= batch_size; ++i) { + cu_seqlens_k_vec.push_back(static_cast(i * kv_seq_len)); + } + } + + auto cu_seqlens_q = + torch::tensor(cu_seqlens_q_vec, torch::dtype(torch::kInt32)) + .to(query.device()) + .contiguous(); + auto cu_seqlens_k = + torch::tensor(cu_seqlens_k_vec, torch::dtype(torch::kInt32)) + .to(query.device()) + .contiguous(); + + auto output = layer::dense_varlen_flash_attention( + dense_query, + dense_key, + dense_value, + cu_seqlens_q, + cu_seqlens_k, + std::pow(static_cast(head_dim), -0.5), + /*is_causal=*/false); + return output.view({batch_size, q_seq_len, num_heads, head_dim}) + .transpose(1, 2) + .contiguous() + .to(output_dtype); +} +#endif + // Global Response Normalization (GRN) for ConvNeXtV2 class AudioGRNImpl : public torch::nn::Module { public: @@ -936,6 +1057,9 @@ class AudioSelfAttentionImpl : public torch::nn::Module { k = apply_rotary_emb(k, rope->first, rope->second); } +#if defined(USE_DCU) + torch::Tensor out = audio_dense_varlen_flash_attention(q, k, v, mask); +#else // Attention mask: convert bool (True=attend) to additive float mask. // cuDNN/Flash backends require a float additive mask, not a bool mask. std::optional attn_mask; @@ -961,6 +1085,7 @@ class AudioSelfAttentionImpl : public torch::nn::Module { torch::Tensor weights = torch::softmax(scores.to(torch::kFloat32), -1).to(q.dtype()); torch::Tensor out = torch::matmul(weights, v); // (B,H,S,head_dim) +#endif out = out.transpose(1, 2).contiguous().view({B, S, inner_dim_}).to(q.dtype()); @@ -1033,6 +1158,9 @@ class AudioCrossAttentionImpl : public torch::nn::Module { k = apply_rotary_emb(k, cond_rope->first, cond_rope->second); } +#if defined(USE_DCU) + torch::Tensor out = audio_dense_varlen_flash_attention(q, k, v, cond_mask); +#else // Build cross-attn mask from cond_mask: (B, H, S, Sc) // cond_mask is True for valid (attend) tokens. // Convert bool to additive float mask for cuDNN/Flash backend @@ -1064,6 +1192,7 @@ class AudioCrossAttentionImpl : public torch::nn::Module { torch::nan_to_num(torch::softmax(scores.to(torch::kFloat32), -1), 0.0) .to(q.dtype()); torch::Tensor out = torch::matmul(weights, v); // (B,H,S,head_dim) +#endif out = out.transpose(1, 2).contiguous().view({B, S, inner_dim_}).to(q.dtype()); diff --git a/xllm/models/models.h b/xllm/models/models.h index 440b31ee23..e976b70880 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -106,7 +106,11 @@ limitations under the License. #include "llm/musa/qwen3.h" // IWYU pragma: keep #elif defined(USE_DCU) #include "dit/pipelines/pipeline_flux.h" +#include "dit/pipelines/pipeline_flux_control.h" +#include "dit/pipelines/pipeline_flux_fill.h" +#include "dit/pipelines/pipeline_longcat_audiodit.h" #include "dit/pipelines/pipeline_longcat_image.h" +#include "dit/pipelines/pipeline_longcat_image_edit.h" #include "dit/pipelines/pipeline_qwenimage_edit_plus.h" #include "dit/pipelines/pipeline_wan_i2v.h" #include "llm/deepseek_v2.h" // IWYU pragma: keep