Skip to content

feat: add DCU support for LongCat and Flux series pipelines. #1937

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
WenQ7 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add DCU support for LongCat and Flux series pipelines. #1937

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions xllm/models/dit/autoencoders/autoencoder_kl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
1 change: 1 addition & 0 deletions xllm/models/dit/pipelines/pipeline_flux.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 6 additions & 0 deletions xllm/models/dit/pipelines/pipeline_flux_control.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
17 changes: 15 additions & 2 deletions xllm/models/dit/pipelines/pipeline_flux_fill.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<torch::Tensor> 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;
Expand All @@ -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,
Expand Down Expand Up @@ -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
24 changes: 17 additions & 7 deletions xllm/models/dit/pipelines/pipeline_longcat_audiodit.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -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());
Expand Down Expand Up @@ -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<DiTModelImpl<LongCatAudioDiTPipeline>>(
std::move(model), context.get_tensor_options());
});
auto factory = [](const DiTModelContext& context) {
LongCatAudioDiTPipeline model(context);
model->eval();
return std::make_unique<DiTModelImpl<LongCatAudioDiTPipeline>>(
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
Expand Down
45 changes: 30 additions & 15 deletions xllm/models/dit/pipelines/pipeline_longcat_image.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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<DiTModelImpl<LongCatImagePipeline>>(
std::move(model), context.get_tensor_options());
});
auto factory = [](const DiTModelContext& context) {
LongCatImagePipeline model(context);
model->eval();
return std::make_unique<DiTModelImpl<LongCatImagePipeline>>(
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
Expand Down
74 changes: 50 additions & 24 deletions xllm/models/dit/pipelines/pipeline_longcat_image_edit.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -80,7 +82,7 @@ inline std::pair<int64_t, int64_t> 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)
Expand All @@ -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();
Expand Down Expand Up @@ -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();
}

Expand Down Expand Up @@ -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<torch::Tensor> 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<torch::Tensor> output = forward_(
image, // input image
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand Down Expand Up @@ -799,6 +820,9 @@ class LongCatImageEditPipelineImpl : public torch::nn::Module {
const torch::Tensor& attention_mask,
std::optional<MMBatchData> 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) {
Expand Down Expand Up @@ -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<DiTModelImpl<LongCatImageEditPipeline>>(
std::move(model), context.get_tensor_options());
});
auto factory = [](const DiTModelContext& context) {
LongCatImageEditPipeline model(context);
model->eval();
return std::make_unique<DiTModelImpl<LongCatImageEditPipeline>>(
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
Expand Down
Loading
Loading