Switch to using DistributedImageModel
This commit is contained in:
ciaranbor
@@ -1,364 +0,0 @@
|
||||
from math import ceil
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
from mflux.callbacks.callbacks import Callbacks
|
||||
from mflux.config.config import Config
|
||||
from mflux.config.model_config import ModelConfig
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator
|
||||
from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
|
||||
from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder
|
||||
from mflux.models.flux.variants.txt2img.flux import Flux1
|
||||
from mflux.utils.array_util import ArrayUtil
|
||||
from mflux.utils.exceptions import StopImageGenerationException
|
||||
from mflux.utils.image_util import ImageUtil
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
|
||||
from exo.shared.types.worker.instances import BoundInstance
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata
|
||||
from exo.worker.download.download_utils import build_model_path
|
||||
from exo.worker.engines.mflux.pipefusion.distributed_denoising import (
|
||||
DistributedDenoising,
|
||||
)
|
||||
from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init, mx_barrier
|
||||
from exo.worker.runner.bootstrap import logger
|
||||
|
||||
|
||||
class DistributedFlux1:
|
||||
"""
|
||||
Wrapper for Flux1 that attaches distributed group and shard metadata.
|
||||
|
||||
This wrapper enables the generation runtime to access distributed context
|
||||
(group, rank, world_size, shard boundaries).
|
||||
"""
|
||||
|
||||
__slots__ = ("_model", "_group", "_shard_metadata", "_variant")
|
||||
|
||||
_model: Flux1
|
||||
_group: Optional[mx.distributed.Group]
|
||||
_shard_metadata: PipelineShardMetadata
|
||||
_variant: str
|
||||
|
||||
_steps_map: dict[str, dict[str, int]] = {
|
||||
"schnell": {"low": 1, "medium": 2, "high": 4},
|
||||
"dev": {"low": 10, "medium": 25, "high": 50},
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_id: str,
|
||||
local_path: Path,
|
||||
shard_metadata: PipelineShardMetadata,
|
||||
group: Optional[mx.distributed.Group] = None,
|
||||
quantize: int | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize DistributedFlux1 directly from configuration.
|
||||
|
||||
Args:
|
||||
model_id: The model identifier (e.g., "black-forest-labs/FLUX.1-schnell")
|
||||
local_path: Path to the local model weights
|
||||
shard_metadata: Pipeline shard metadata with layer assignments
|
||||
group: MLX distributed group for multi-node coordination (None for single-node)
|
||||
quantize: Optional quantization bit width
|
||||
"""
|
||||
variant = model_id.split("-")[-1]
|
||||
model = Flux1(
|
||||
model_config=ModelConfig.from_name(model_name=model_id, base_model=None),
|
||||
local_path=str(local_path),
|
||||
quantize=quantize,
|
||||
)
|
||||
|
||||
if group is not None:
|
||||
# Apply pipeline parallelism by wrapping the transformer
|
||||
model.transformer = DistributedDenoising( # type: ignore[assignment]
|
||||
transformer=model.transformer,
|
||||
group=group,
|
||||
shard_metadata=shard_metadata,
|
||||
num_sync_steps=ceil(self._steps_map[variant]["medium"] / 8),
|
||||
)
|
||||
logger.info("Applied pipefusion transformations")
|
||||
|
||||
mx.eval(model.parameters())
|
||||
|
||||
# TODO: Do we need this?
|
||||
mx.eval(model)
|
||||
|
||||
# Synchronize processes before generation to avoid timeout
|
||||
mx_barrier(group)
|
||||
logger.info(f"Flux transformer sharded for rank {group.rank()}")
|
||||
else:
|
||||
logger.info("Single-node Flux initialization")
|
||||
|
||||
object.__setattr__(self, "_model", model)
|
||||
object.__setattr__(self, "_group", group)
|
||||
object.__setattr__(self, "_shard_metadata", shard_metadata)
|
||||
object.__setattr__(self, "_variant", variant)
|
||||
|
||||
@classmethod
|
||||
def from_bound_instance(cls, bound_instance: BoundInstance) -> "DistributedFlux1":
|
||||
"""
|
||||
Create DistributedFlux1 from a BoundInstance.
|
||||
|
||||
This factory method extracts model configuration from the bound instance
|
||||
and handles distributed initialization if needed.
|
||||
|
||||
Args:
|
||||
bound_instance: The bound instance containing model and shard info
|
||||
|
||||
Returns:
|
||||
Initialized DistributedFlux1 ready for inference
|
||||
"""
|
||||
model_id = bound_instance.bound_shard.model_meta.model_id
|
||||
model_path = build_model_path(model_id)
|
||||
|
||||
shard_metadata = bound_instance.bound_shard
|
||||
if not isinstance(shard_metadata, PipelineShardMetadata):
|
||||
raise ValueError("Expected PipelineShardMetadata for Flux")
|
||||
|
||||
is_distributed = (
|
||||
len(bound_instance.instance.shard_assignments.node_to_runner) > 1
|
||||
)
|
||||
|
||||
if is_distributed:
|
||||
logger.info("Starting distributed init for Flux")
|
||||
group = mlx_distributed_init(bound_instance)
|
||||
else:
|
||||
group = None
|
||||
|
||||
return cls(
|
||||
model_id=model_id,
|
||||
local_path=model_path,
|
||||
shard_metadata=shard_metadata,
|
||||
group=group,
|
||||
)
|
||||
|
||||
@property
|
||||
def model(self) -> Flux1:
|
||||
"""The underlying Flux1 model."""
|
||||
return self._model
|
||||
|
||||
@property
|
||||
def group(self) -> Optional[mx.distributed.Group]:
|
||||
"""The MLX distributed group for this model (None for single-node)."""
|
||||
return self._group
|
||||
|
||||
@property
|
||||
def shard_metadata(self) -> PipelineShardMetadata:
|
||||
"""Shard metadata containing layer assignments and device info."""
|
||||
return self._shard_metadata
|
||||
|
||||
@property
|
||||
def rank(self) -> int:
|
||||
"""This device's rank in the distributed group."""
|
||||
return self._shard_metadata.device_rank
|
||||
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
"""Total number of devices in the distributed group."""
|
||||
return self._shard_metadata.world_size
|
||||
|
||||
@property
|
||||
def is_first_stage(self) -> bool:
|
||||
"""True if this device is the first stage in the pipeline."""
|
||||
return self._shard_metadata.device_rank == 0
|
||||
|
||||
@property
|
||||
def is_last_stage(self) -> bool:
|
||||
"""True if this device is the last stage in the pipeline."""
|
||||
return self._shard_metadata.device_rank == self._shard_metadata.world_size - 1
|
||||
|
||||
@property
|
||||
def is_distributed(self) -> bool:
|
||||
"""True if running in distributed mode (world_size > 1)."""
|
||||
return self._shard_metadata.world_size > 1
|
||||
|
||||
# Delegate attribute access to the underlying model.
|
||||
# Guarded with TYPE_CHECKING to prevent type checker complaints
|
||||
# while still providing full delegation at runtime.
|
||||
if not TYPE_CHECKING:
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._model, name)
|
||||
|
||||
def __setattr__(self, name: str, value: Any) -> None:
|
||||
if name in ("_model", "_group", "_shard_metadata"):
|
||||
object.__setattr__(self, name, value)
|
||||
else:
|
||||
setattr(self._model, name, value)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
height: int,
|
||||
width: int,
|
||||
quality: Literal["low", "medium", "high"] = "medium",
|
||||
seed: int = 2,
|
||||
) -> Optional[Image.Image]:
|
||||
"""
|
||||
Generate an image using the Flux1 model.
|
||||
|
||||
For distributed inference, only the first stage (rank 0) returns the image.
|
||||
Other stages return None after participating in the pipeline.
|
||||
|
||||
Args:
|
||||
prompt: Text description of the image to generate
|
||||
height: Image height in pixels
|
||||
width: Image width in pixels
|
||||
quality: Generation quality ("low", "medium", "high")
|
||||
seed: Random seed for reproducibility
|
||||
|
||||
Returns:
|
||||
Generated PIL Image (rank 0) or None (other ranks)
|
||||
"""
|
||||
# Determine number of inference steps based on quality
|
||||
steps = self._steps_map[self._variant][quality]
|
||||
|
||||
config = Config(num_inference_steps=steps, height=height, width=width)
|
||||
image = self._generate_image(settings=config, prompt=prompt, seed=seed)
|
||||
logger.info("generated image")
|
||||
|
||||
# Only rank 0 returns the actual image
|
||||
if self.is_first_stage:
|
||||
return image.image
|
||||
|
||||
def _generate_image(self, settings: Config, prompt: str, seed: int) -> Any:
|
||||
"""
|
||||
Internal image generation with the diffusion loop.
|
||||
|
||||
This method implements the core diffusion process with distributed
|
||||
communication handled at the transformer level (via block wrappers
|
||||
for now, will be moved to loop level for async pipeline).
|
||||
"""
|
||||
model = self._model
|
||||
|
||||
# 0. Create runtime config
|
||||
config = RuntimeConfig(settings, model.model_config)
|
||||
time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
|
||||
|
||||
# 1. Create initial latents (all nodes create the same latents with same seed)
|
||||
latents = LatentCreator.create_for_txt2img_or_img2img(
|
||||
seed=seed,
|
||||
height=config.height,
|
||||
width=config.width,
|
||||
img2img=Img2Img(
|
||||
vae=model.vae,
|
||||
latent_creator=FluxLatentCreator,
|
||||
image_path=config.image_path,
|
||||
sigmas=config.scheduler.sigmas,
|
||||
init_time_step=config.init_time_step,
|
||||
),
|
||||
)
|
||||
|
||||
# 2. Encode the prompt (all nodes encode to get consistent embeddings)
|
||||
prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_cache=model.prompt_cache,
|
||||
t5_tokenizer=model.t5_tokenizer,
|
||||
clip_tokenizer=model.clip_tokenizer,
|
||||
t5_text_encoder=model.t5_text_encoder,
|
||||
clip_text_encoder=model.clip_text_encoder,
|
||||
)
|
||||
|
||||
# (Optional) Call subscribers for beginning of loop
|
||||
Callbacks.before_loop(
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# 3. Main diffusion loop
|
||||
for t in time_steps:
|
||||
try:
|
||||
latents = self._diffusion_step(
|
||||
t=t,
|
||||
config=config,
|
||||
latents=latents,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
)
|
||||
|
||||
# (Optional) Call subscribers in-loop
|
||||
Callbacks.in_loop(
|
||||
t=t,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
time_steps=time_steps,
|
||||
)
|
||||
|
||||
mx.eval(latents)
|
||||
|
||||
except KeyboardInterrupt: # noqa: PERF203
|
||||
Callbacks.interruption(
|
||||
t=t,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
time_steps=time_steps,
|
||||
)
|
||||
raise StopImageGenerationException(
|
||||
f"Stopping image generation at step {t + 1}/{len(time_steps)}"
|
||||
) from None
|
||||
|
||||
# (Optional) Call subscribers after loop
|
||||
Callbacks.after_loop(
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# 4. Decode latents to image (all nodes decode for now)
|
||||
latents = ArrayUtil.unpack_latents(
|
||||
latents=latents, height=config.height, width=config.width
|
||||
)
|
||||
decoded = model.vae.decode(latents)
|
||||
return ImageUtil.to_image(
|
||||
decoded_latents=decoded,
|
||||
config=config,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
quantization=model.bits,
|
||||
lora_paths=model.lora_paths,
|
||||
lora_scales=model.lora_scales,
|
||||
image_path=config.image_path,
|
||||
image_strength=config.image_strength,
|
||||
generation_time=time_steps.format_dict["elapsed"],
|
||||
)
|
||||
|
||||
def _diffusion_step(
|
||||
self,
|
||||
t: int,
|
||||
config: RuntimeConfig,
|
||||
latents: mx.array,
|
||||
prompt_embeds: mx.array,
|
||||
pooled_prompt_embeds: mx.array,
|
||||
) -> mx.array:
|
||||
if self._group is None:
|
||||
latents = config.scheduler.scale_model_input(latents, t)
|
||||
|
||||
noise = self._model.transformer(
|
||||
t=t,
|
||||
config=config,
|
||||
hidden_states=latents,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
)
|
||||
|
||||
if self._group is None:
|
||||
latents = config.scheduler.step(
|
||||
model_output=noise,
|
||||
timestep=t,
|
||||
sample=latents,
|
||||
)
|
||||
else:
|
||||
latents = noise
|
||||
|
||||
return latents
|
||||
@@ -6,7 +6,7 @@ from PIL import Image
|
||||
|
||||
from exo.shared.types.api import ImageGenerationTaskParams
|
||||
from exo.shared.types.worker.runner_response import ImageGenerationResponse
|
||||
from exo.worker.engines.mflux.distributed_flux import DistributedFlux1
|
||||
from exo.worker.engines.mflux.distributed_model import DistributedImageModel
|
||||
|
||||
image_generation_stream = mx.new_stream(mx.default_device())
|
||||
|
||||
@@ -30,7 +30,7 @@ def parse_size(size_str: str | None) -> tuple[int, int]:
|
||||
return (1024, 1024)
|
||||
|
||||
|
||||
def warmup_mflux(model: DistributedFlux1) -> Image.Image | None:
|
||||
def warmup_mflux(model: DistributedImageModel) -> Image.Image | None:
|
||||
"""
|
||||
Warmup the model by generating a small image.
|
||||
|
||||
@@ -46,11 +46,11 @@ def warmup_mflux(model: DistributedFlux1) -> Image.Image | None:
|
||||
|
||||
|
||||
def mflux_generate(
|
||||
model: DistributedFlux1,
|
||||
model: DistributedImageModel,
|
||||
task: ImageGenerationTaskParams,
|
||||
) -> Generator[ImageGenerationResponse, None, None]:
|
||||
"""
|
||||
Generate an image using the DistributedFlux1 model.
|
||||
Generate an image using the DistributedImageModel.
|
||||
|
||||
For distributed inference, only rank 0 yields the response.
|
||||
Other ranks participate in the pipeline but yield nothing.
|
||||
|
||||
@@ -171,11 +171,11 @@ class DistributedDenoising:
|
||||
def _create_patches(
|
||||
self,
|
||||
latents: mx.array,
|
||||
runtime_config: RuntimeConfig,
|
||||
config: RuntimeConfig,
|
||||
) -> tuple[list[mx.array], list[tuple[int, int]]]:
|
||||
# Calculate patch metadata using model config
|
||||
latent_height = runtime_config.height // self.config.vae_scale_factor
|
||||
latent_width = runtime_config.width // self.config.vae_scale_factor
|
||||
latent_height = config.height // self.config.vae_scale_factor
|
||||
latent_width = config.width // self.config.vae_scale_factor
|
||||
patch_size = self.config.patch_size
|
||||
|
||||
patch_heights, _ = calculate_patch_heights(
|
||||
@@ -191,7 +191,7 @@ class DistributedDenoising:
|
||||
def _sync_pipeline(
|
||||
self,
|
||||
t: int,
|
||||
runtime_config: RuntimeConfig,
|
||||
config: RuntimeConfig,
|
||||
hidden_states: mx.array,
|
||||
prompt_embeds: mx.array,
|
||||
pooled_prompt_embeds: mx.array,
|
||||
@@ -199,7 +199,7 @@ class DistributedDenoising:
|
||||
) -> mx.array:
|
||||
prev_latents = hidden_states
|
||||
|
||||
hidden_states = runtime_config.scheduler.scale_model_input(hidden_states, t)
|
||||
hidden_states = config.scheduler.scale_model_input(hidden_states, t)
|
||||
|
||||
# === PHASE 1: Embeddings ===
|
||||
# First stage: compute embeddings
|
||||
@@ -211,12 +211,12 @@ class DistributedDenoising:
|
||||
|
||||
# All stages need these for their blocks
|
||||
text_embeddings = self.adapter.compute_text_embeddings(
|
||||
t, pooled_prompt_embeds, self.transformer, runtime_config
|
||||
t, pooled_prompt_embeds, self.transformer, config
|
||||
)
|
||||
image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
|
||||
prompt_embeds,
|
||||
self.transformer,
|
||||
runtime_config,
|
||||
config,
|
||||
kontext_image_ids=kontext_image_ids,
|
||||
)
|
||||
|
||||
@@ -335,7 +335,7 @@ class DistributedDenoising:
|
||||
hidden_states, text_embeddings, self.transformer
|
||||
)
|
||||
|
||||
hidden_states = runtime_config.scheduler.step(
|
||||
hidden_states = config.scheduler.step(
|
||||
model_output=hidden_states,
|
||||
timestep=t,
|
||||
sample=prev_latents,
|
||||
@@ -360,7 +360,7 @@ class DistributedDenoising:
|
||||
def _async_pipeline(
|
||||
self,
|
||||
t: int,
|
||||
runtime_config: RuntimeConfig,
|
||||
config: RuntimeConfig,
|
||||
patch_latents: list[mx.array],
|
||||
token_indices: list[tuple[int, int]],
|
||||
prompt_embeds: mx.array,
|
||||
@@ -371,15 +371,15 @@ class DistributedDenoising:
|
||||
assert self.single_kv_caches is not None
|
||||
|
||||
# TODO(ciaran): needed in general?
|
||||
# hidden_states = runtime_config.scheduler.scale_model_input(hidden_states, t)
|
||||
# hidden_states = config.scheduler.scale_model_input(hidden_states, t)
|
||||
|
||||
text_embeddings = self.adapter.compute_text_embeddings(
|
||||
t, pooled_prompt_embeds, self.transformer, runtime_config
|
||||
t, pooled_prompt_embeds, self.transformer, config
|
||||
)
|
||||
image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
|
||||
prompt_embeds,
|
||||
self.transformer,
|
||||
runtime_config,
|
||||
config,
|
||||
kontext_image_ids=kontext_image_ids,
|
||||
)
|
||||
|
||||
@@ -503,7 +503,7 @@ class DistributedDenoising:
|
||||
patch_img_only, text_embeddings, self.transformer
|
||||
)
|
||||
|
||||
patch = runtime_config.scheduler.step(
|
||||
patch = config.scheduler.step(
|
||||
model_output=patch_img_only,
|
||||
timestep=t,
|
||||
sample=patch_prev,
|
||||
@@ -521,7 +521,7 @@ class DistributedDenoising:
|
||||
def __call__(
|
||||
self,
|
||||
t: int,
|
||||
runtime_config: RuntimeConfig,
|
||||
config: RuntimeConfig,
|
||||
hidden_states: mx.array,
|
||||
prompt_embeds: mx.array,
|
||||
pooled_prompt_embeds: mx.array,
|
||||
@@ -532,20 +532,18 @@ class DistributedDenoising:
|
||||
if t < self.num_sync_steps:
|
||||
latents = self._sync_pipeline(
|
||||
t,
|
||||
runtime_config,
|
||||
config,
|
||||
hidden_states,
|
||||
prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
kontext_image_ids,
|
||||
)
|
||||
else:
|
||||
patch_latents, token_indices = self._create_patches(
|
||||
hidden_states, runtime_config
|
||||
)
|
||||
patch_latents, token_indices = self._create_patches(hidden_states, config)
|
||||
|
||||
patch_latents = self._async_pipeline(
|
||||
t,
|
||||
runtime_config,
|
||||
config,
|
||||
patch_latents,
|
||||
token_indices,
|
||||
prompt_embeds,
|
||||
@@ -555,7 +553,7 @@ class DistributedDenoising:
|
||||
|
||||
# Receive final patches from last rank
|
||||
if (
|
||||
t == runtime_config.num_inference_steps - 1
|
||||
t == config.num_inference_steps - 1
|
||||
and self.is_first_stage
|
||||
and not self.is_last_stage
|
||||
):
|
||||
|
||||
@@ -1,644 +0,0 @@
|
||||
import mlx.core as mx
|
||||
from mflux.models.flux.model.flux_transformer.common.attention_utils import (
|
||||
AttentionUtils,
|
||||
)
|
||||
from mflux.models.flux.model.flux_transformer.joint_attention import JointAttention
|
||||
from mflux.models.flux.model.flux_transformer.joint_transformer_block import (
|
||||
JointTransformerBlock,
|
||||
)
|
||||
from mflux.models.flux.model.flux_transformer.single_block_attention import (
|
||||
SingleBlockAttention,
|
||||
)
|
||||
from mflux.models.flux.model.flux_transformer.single_transformer_block import (
|
||||
SingleTransformerBlock,
|
||||
)
|
||||
|
||||
from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
|
||||
|
||||
|
||||
class CachedJointAttention:
|
||||
"""JointAttention that uses KV cache for PipeFusion.
|
||||
|
||||
Processes attention for a single image patch while attending to the full
|
||||
sequence (text + all image patches) using cached K/V values.
|
||||
"""
|
||||
|
||||
head_dim = 128
|
||||
num_heads = 24
|
||||
|
||||
def __init__(self, attn: JointAttention):
|
||||
"""Wrap an existing JointAttention module.
|
||||
|
||||
Args:
|
||||
attn: The original JointAttention module (for weight access)
|
||||
"""
|
||||
self.attn = attn
|
||||
self.num_heads = attn.num_heads
|
||||
self.head_dim = attn.head_dimension
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
norm_hidden: mx.array,
|
||||
norm_encoder: mx.array,
|
||||
image_rotary_emb: mx.array,
|
||||
kv_cache: ImagePatchKVCache,
|
||||
patch_start: int,
|
||||
patch_end: int,
|
||||
text_seq_len: int,
|
||||
) -> tuple[mx.array, mx.array]:
|
||||
"""Forward pass with KV caching for patch-based attention.
|
||||
|
||||
Args:
|
||||
norm_hidden: Normalized image patch hidden states [B, patch_img_len, D]
|
||||
norm_encoder: Normalized full text hidden states [B, text_len, D]
|
||||
image_rotary_emb: Full rotary embeddings for [text + full_image]
|
||||
kv_cache: KV cache to update and read from
|
||||
patch_start: Start token index of this patch in image sequence
|
||||
patch_end: End token index of this patch in image sequence
|
||||
text_seq_len: Length of text sequence
|
||||
|
||||
Returns:
|
||||
Tuple of (hidden_states_out, encoder_hidden_states_out) for the patch
|
||||
"""
|
||||
# 1. Compute Q, K, V for current image patch
|
||||
img_query, img_key, img_value = AttentionUtils.process_qkv(
|
||||
hidden_states=norm_hidden,
|
||||
to_q=self.attn.to_q,
|
||||
to_k=self.attn.to_k,
|
||||
to_v=self.attn.to_v,
|
||||
norm_q=self.attn.norm_q,
|
||||
norm_k=self.attn.norm_k,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 2. Compute Q, K, V for text (always full)
|
||||
txt_query, txt_key, txt_value = AttentionUtils.process_qkv(
|
||||
hidden_states=norm_encoder,
|
||||
to_q=self.attn.add_q_proj,
|
||||
to_k=self.attn.add_k_proj,
|
||||
to_v=self.attn.add_v_proj,
|
||||
norm_q=self.attn.norm_added_q,
|
||||
norm_k=self.attn.norm_added_k,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 3. Concatenate Q for attention output: [text, image_patch]
|
||||
query = mx.concatenate([txt_query, img_query], axis=2)
|
||||
|
||||
# 4. Concatenate K, V for this patch: [text, image_patch]
|
||||
patch_key = mx.concatenate([txt_key, img_key], axis=2)
|
||||
patch_value = mx.concatenate([txt_value, img_value], axis=2)
|
||||
|
||||
# 5. Extract RoPE for [text + current_patch] positions
|
||||
# RoPE shape is (1, 1, seq_len, head_dim/2, 2, 2), sequence is on axis 2
|
||||
text_rope = image_rotary_emb[:, :, :text_seq_len, ...]
|
||||
patch_img_rope = image_rotary_emb[
|
||||
:, :, text_seq_len + patch_start : text_seq_len + patch_end, ...
|
||||
]
|
||||
patch_rope = mx.concatenate([text_rope, patch_img_rope], axis=2)
|
||||
|
||||
# 6. Apply RoPE to Q and K for current patch
|
||||
query, patch_key = AttentionUtils.apply_rope(
|
||||
xq=query, xk=patch_key, freqs_cis=patch_rope
|
||||
)
|
||||
|
||||
# 7. Update cache with this patch's IMAGE K/V only (after RoPE)
|
||||
# Text K/V is not cached - it's always fresh and the same for all patches
|
||||
kv_cache.update_image_patch(
|
||||
patch_start=patch_start,
|
||||
patch_end=patch_end,
|
||||
key=patch_key[:, :, text_seq_len:, :],
|
||||
value=patch_value[:, :, text_seq_len:, :],
|
||||
)
|
||||
|
||||
# 8. Get full K, V by concatenating fresh text K/V with cached image K/V
|
||||
# Text K/V: fresh (just computed), Image K/V: fresh for current patch, stale for others
|
||||
full_key, full_value = kv_cache.get_full_kv(
|
||||
text_key=patch_key[:, :, :text_seq_len, :],
|
||||
text_value=patch_value[:, :, :text_seq_len, :],
|
||||
)
|
||||
|
||||
# 9. Compute attention: patch query attends to full K, V
|
||||
# Query shape: [B, H, text_seq_len + patch_len, D]
|
||||
# Key/Value shape: [B, H, text_seq_len + full_img_len, D]
|
||||
# Output shape: [B, text_seq_len + patch_len, D]
|
||||
batch_size = norm_hidden.shape[0]
|
||||
attn_output = AttentionUtils.compute_attention(
|
||||
query=query,
|
||||
key=full_key,
|
||||
value=full_value,
|
||||
batch_size=batch_size,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 10. Extract outputs for text and image patch
|
||||
encoder_hidden_states_out = attn_output[:, :text_seq_len, :]
|
||||
hidden_states_out = attn_output[:, text_seq_len:, :]
|
||||
|
||||
# 11. Project outputs
|
||||
hidden_states_out = self.attn.to_out[0](hidden_states_out)
|
||||
encoder_hidden_states_out = self.attn.to_add_out(encoder_hidden_states_out)
|
||||
|
||||
return hidden_states_out, encoder_hidden_states_out
|
||||
|
||||
|
||||
class CachedSingleBlockAttention:
|
||||
"""SingleBlockAttention that uses KV cache for PipeFusion.
|
||||
|
||||
Processes attention for a single patch of concatenated [text + image]
|
||||
while attending to the full sequence using cached K/V values.
|
||||
"""
|
||||
|
||||
head_dim = 128
|
||||
num_heads = 24
|
||||
|
||||
def __init__(self, attn: SingleBlockAttention):
|
||||
"""Wrap an existing SingleBlockAttention module.
|
||||
|
||||
Args:
|
||||
attn: The original SingleBlockAttention module (for weight access)
|
||||
"""
|
||||
self.attn = attn
|
||||
self.head_dim = attn.head_dimension
|
||||
self.num_heads = attn.num_heads
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
norm_hidden: mx.array,
|
||||
image_rotary_emb: mx.array,
|
||||
kv_cache: ImagePatchKVCache,
|
||||
patch_start: int,
|
||||
patch_end: int,
|
||||
text_seq_len: int,
|
||||
) -> mx.array:
|
||||
"""Forward pass with KV caching for patch-based attention.
|
||||
|
||||
Args:
|
||||
norm_hidden: Normalized patch of [text + image] hidden states [B, text_len + patch_img_len, D]
|
||||
image_rotary_emb: Full rotary embeddings for [text + full_image]
|
||||
kv_cache: KV cache to update and read from
|
||||
patch_start: Start token index of image patch (within image portion)
|
||||
patch_end: End token index of image patch (within image portion)
|
||||
text_seq_len: Length of text portion in hidden_states
|
||||
|
||||
Returns:
|
||||
Attention output for the patch [B, text_len + patch_img_len, D]
|
||||
"""
|
||||
# 1. Compute Q, K, V for current patch
|
||||
query, key, value = AttentionUtils.process_qkv(
|
||||
hidden_states=norm_hidden,
|
||||
to_q=self.attn.to_q,
|
||||
to_k=self.attn.to_k,
|
||||
to_v=self.attn.to_v,
|
||||
norm_q=self.attn.norm_q,
|
||||
norm_k=self.attn.norm_k,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 2. Extract RoPE for [text + current_image_patch] positions
|
||||
# RoPE shape is (1, 1, seq_len, head_dim/2, 2, 2), sequence is on axis 2
|
||||
text_rope = image_rotary_emb[:, :, :text_seq_len, ...]
|
||||
patch_img_rope = image_rotary_emb[
|
||||
:, :, text_seq_len + patch_start : text_seq_len + patch_end, ...
|
||||
]
|
||||
patch_rope = mx.concatenate([text_rope, patch_img_rope], axis=2)
|
||||
|
||||
# 3. Apply RoPE to Q and K
|
||||
query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=patch_rope)
|
||||
|
||||
# 4. Update cache with this patch's IMAGE K/V only (after RoPE)
|
||||
# Text K/V is not cached - it's always fresh and the same for all patches
|
||||
kv_cache.update_image_patch(
|
||||
patch_start=patch_start,
|
||||
patch_end=patch_end,
|
||||
key=key[:, :, text_seq_len:, :],
|
||||
value=value[:, :, text_seq_len:, :],
|
||||
)
|
||||
|
||||
# 5. Get full K, V by concatenating fresh text K/V with cached image K/V
|
||||
# Text K/V: fresh (just computed), Image K/V: fresh for current patch, stale for others
|
||||
full_key, full_value = kv_cache.get_full_kv(
|
||||
text_key=key[:, :, :text_seq_len, :],
|
||||
text_value=value[:, :, :text_seq_len, :],
|
||||
)
|
||||
|
||||
# 6. Compute attention: patch query attends to full K, V
|
||||
batch_size = norm_hidden.shape[0]
|
||||
attn_output = AttentionUtils.compute_attention(
|
||||
query=query,
|
||||
key=full_key,
|
||||
value=full_value,
|
||||
batch_size=batch_size,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
return attn_output
|
||||
|
||||
|
||||
class PatchedJointTransformerBlock:
|
||||
"""Joint transformer block with KV caching for patch-based processing.
|
||||
|
||||
Wraps a JointTransformerBlock to process image patches while using
|
||||
cached K/V values for the full sequence (text + all image patches).
|
||||
"""
|
||||
|
||||
def __init__(self, block: JointTransformerBlock):
|
||||
"""Wrap an existing JointTransformerBlock.
|
||||
|
||||
Args:
|
||||
block: The original JointTransformerBlock (for weight access)
|
||||
"""
|
||||
self.block = block
|
||||
self.cached_attn = CachedJointAttention(block.attn)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
patch_hidden: mx.array,
|
||||
encoder_hidden_states: mx.array,
|
||||
text_embeddings: mx.array,
|
||||
image_rotary_emb: mx.array,
|
||||
kv_cache: ImagePatchKVCache,
|
||||
patch_start: int,
|
||||
patch_end: int,
|
||||
text_seq_len: int,
|
||||
) -> tuple[mx.array, mx.array]:
|
||||
"""Forward pass with KV caching for patch-based processing.
|
||||
|
||||
Args:
|
||||
patch_hidden: Image patch hidden states [B, patch_img_len, D]
|
||||
encoder_hidden_states: Full text hidden states [B, text_len, D]
|
||||
text_embeddings: Time + pooled text conditioning
|
||||
image_rotary_emb: Full rotary embeddings for [text + full_image]
|
||||
kv_cache: KV cache to update and read from
|
||||
patch_start: Start token index of this patch in image sequence
|
||||
patch_end: End token index of this patch in image sequence
|
||||
text_seq_len: Length of text sequence
|
||||
|
||||
Returns:
|
||||
Tuple of (encoder_hidden_states, patch_hidden) after block processing
|
||||
"""
|
||||
# 1. Compute norms
|
||||
norm_hidden, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.block.norm1(
|
||||
hidden_states=patch_hidden,
|
||||
text_embeddings=text_embeddings,
|
||||
)
|
||||
(
|
||||
norm_encoder,
|
||||
c_gate_msa,
|
||||
c_shift_mlp,
|
||||
c_scale_mlp,
|
||||
c_gate_mlp,
|
||||
) = self.block.norm1_context(
|
||||
hidden_states=encoder_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
)
|
||||
|
||||
# 2. Compute attention with KV cache
|
||||
attn_output, context_attn_output = self.cached_attn(
|
||||
norm_hidden=norm_hidden,
|
||||
norm_encoder=norm_encoder,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
kv_cache=kv_cache,
|
||||
patch_start=patch_start,
|
||||
patch_end=patch_end,
|
||||
text_seq_len=text_seq_len,
|
||||
)
|
||||
|
||||
# 3. Apply norm and feed forward
|
||||
patch_hidden = JointTransformerBlock.apply_norm_and_feed_forward(
|
||||
hidden_states=patch_hidden,
|
||||
attn_output=attn_output,
|
||||
gate_mlp=gate_mlp,
|
||||
gate_msa=gate_msa,
|
||||
scale_mlp=scale_mlp,
|
||||
shift_mlp=shift_mlp,
|
||||
norm_layer=self.block.norm2,
|
||||
ff_layer=self.block.ff,
|
||||
)
|
||||
encoder_hidden_states = JointTransformerBlock.apply_norm_and_feed_forward(
|
||||
hidden_states=encoder_hidden_states,
|
||||
attn_output=context_attn_output,
|
||||
gate_mlp=c_gate_mlp,
|
||||
gate_msa=c_gate_msa,
|
||||
scale_mlp=c_scale_mlp,
|
||||
shift_mlp=c_shift_mlp,
|
||||
norm_layer=self.block.norm2_context,
|
||||
ff_layer=self.block.ff_context,
|
||||
)
|
||||
|
||||
return encoder_hidden_states, patch_hidden
|
||||
|
||||
|
||||
class PatchedSingleTransformerBlock:
|
||||
"""Single transformer block with KV caching for patch-based processing.
|
||||
|
||||
Wraps a SingleTransformerBlock to process patches of concatenated [text + image]
|
||||
while using cached K/V values for the full sequence.
|
||||
"""
|
||||
|
||||
def __init__(self, block: SingleTransformerBlock):
|
||||
"""Wrap an existing SingleTransformerBlock.
|
||||
|
||||
Args:
|
||||
block: The original SingleTransformerBlock (for weight access)
|
||||
"""
|
||||
self.block = block
|
||||
self.cached_attn = CachedSingleBlockAttention(block.attn)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
patch_hidden: mx.array,
|
||||
text_embeddings: mx.array,
|
||||
image_rotary_emb: mx.array,
|
||||
kv_cache: ImagePatchKVCache,
|
||||
patch_start: int,
|
||||
patch_end: int,
|
||||
text_seq_len: int,
|
||||
) -> mx.array:
|
||||
"""Forward pass with KV caching for patch-based processing.
|
||||
|
||||
Args:
|
||||
patch_hidden: Patch of [text + image] hidden states [B, text_len + patch_img_len, D]
|
||||
text_embeddings: Time + pooled text conditioning
|
||||
image_rotary_emb: Full rotary embeddings for [text + full_image]
|
||||
kv_cache: KV cache to update and read from
|
||||
patch_start: Start token index of image patch (within image portion)
|
||||
patch_end: End token index of image patch (within image portion)
|
||||
text_seq_len: Length of text portion in hidden_states
|
||||
|
||||
Returns:
|
||||
Output hidden states [B, text_len + patch_img_len, D]
|
||||
"""
|
||||
# 0. Establish residual connection
|
||||
residual = patch_hidden
|
||||
|
||||
# 1. Compute norm
|
||||
norm_hidden, gate = self.block.norm(
|
||||
hidden_states=patch_hidden,
|
||||
text_embeddings=text_embeddings,
|
||||
)
|
||||
|
||||
# 2. Compute attention with KV cache
|
||||
attn_output = self.cached_attn(
|
||||
norm_hidden=norm_hidden,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
kv_cache=kv_cache,
|
||||
patch_start=patch_start,
|
||||
patch_end=patch_end,
|
||||
text_seq_len=text_seq_len,
|
||||
)
|
||||
|
||||
# 3. Apply feed forward and projection
|
||||
hidden_states = self.block._apply_feed_forward_and_projection(
|
||||
norm_hidden_states=norm_hidden,
|
||||
attn_output=attn_output,
|
||||
gate=gate,
|
||||
)
|
||||
|
||||
return residual + hidden_states
|
||||
|
||||
|
||||
class CachingJointTransformerBlock:
|
||||
"""Joint transformer block that captures K/V for cache during sync mode.
|
||||
|
||||
Implements the full block computation while storing K/V in the cache for
|
||||
subsequent async timesteps to use as stale values. Computes K/V only once.
|
||||
"""
|
||||
|
||||
def __init__(self, block: JointTransformerBlock, kv_cache: ImagePatchKVCache):
|
||||
"""Wrap an existing JointTransformerBlock with a KV cache.
|
||||
|
||||
Args:
|
||||
block: The original JointTransformerBlock
|
||||
kv_cache: KV cache to populate during forward pass
|
||||
"""
|
||||
self.block = block
|
||||
self.kv_cache = kv_cache
|
||||
self.attn = block.attn
|
||||
self.num_heads = block.attn.num_heads
|
||||
self.head_dim = block.attn.head_dimension
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
encoder_hidden_states: mx.array,
|
||||
text_embeddings: mx.array,
|
||||
rotary_embeddings: mx.array,
|
||||
) -> tuple[mx.array, mx.array]:
|
||||
"""Forward pass that computes attention and populates the KV cache.
|
||||
|
||||
Args:
|
||||
hidden_states: Full image hidden states [B, img_len, D]
|
||||
encoder_hidden_states: Full text hidden states [B, text_len, D]
|
||||
text_embeddings: Time + pooled text conditioning
|
||||
rotary_embeddings: Full rotary embeddings for [text + full_image]
|
||||
|
||||
Returns:
|
||||
Tuple of (encoder_hidden_states, hidden_states) after block processing
|
||||
"""
|
||||
text_seq_len = encoder_hidden_states.shape[1]
|
||||
num_img_tokens = hidden_states.shape[1]
|
||||
batch_size = hidden_states.shape[0]
|
||||
|
||||
# 1. Compute norms (same as PatchedJointTransformerBlock)
|
||||
norm_hidden, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.block.norm1(
|
||||
hidden_states=hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
)
|
||||
(
|
||||
norm_encoder,
|
||||
c_gate_msa,
|
||||
c_shift_mlp,
|
||||
c_scale_mlp,
|
||||
c_gate_mlp,
|
||||
) = self.block.norm1_context(
|
||||
hidden_states=encoder_hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
)
|
||||
|
||||
# 2. Compute Q, K, V for full image
|
||||
img_query, img_key, img_value = AttentionUtils.process_qkv(
|
||||
hidden_states=norm_hidden,
|
||||
to_q=self.attn.to_q,
|
||||
to_k=self.attn.to_k,
|
||||
to_v=self.attn.to_v,
|
||||
norm_q=self.attn.norm_q,
|
||||
norm_k=self.attn.norm_k,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 3. Compute Q, K, V for text
|
||||
txt_query, txt_key, txt_value = AttentionUtils.process_qkv(
|
||||
hidden_states=norm_encoder,
|
||||
to_q=self.attn.add_q_proj,
|
||||
to_k=self.attn.add_k_proj,
|
||||
to_v=self.attn.add_v_proj,
|
||||
norm_q=self.attn.norm_added_q,
|
||||
norm_k=self.attn.norm_added_k,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 4. Concatenate Q, K, V: [text, image]
|
||||
query = mx.concatenate([txt_query, img_query], axis=2)
|
||||
key = mx.concatenate([txt_key, img_key], axis=2)
|
||||
value = mx.concatenate([txt_value, img_value], axis=2)
|
||||
|
||||
# 5. Apply RoPE to Q and K
|
||||
query, key = AttentionUtils.apply_rope(
|
||||
xq=query, xk=key, freqs_cis=rotary_embeddings
|
||||
)
|
||||
|
||||
# 6. Store only IMAGE K/V in cache for async pipeline warmstart
|
||||
# Text K/V is not cached - it's always computed fresh
|
||||
self.kv_cache.update_image_patch(
|
||||
patch_start=0,
|
||||
patch_end=num_img_tokens,
|
||||
key=key[:, :, text_seq_len:, :],
|
||||
value=value[:, :, text_seq_len:, :],
|
||||
)
|
||||
|
||||
# 7. Compute full attention (standard, not using stale cache)
|
||||
attn_output = AttentionUtils.compute_attention(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
batch_size=batch_size,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 8. Extract outputs for text and image
|
||||
context_attn_output = attn_output[:, :text_seq_len, :]
|
||||
attn_output = attn_output[:, text_seq_len:, :]
|
||||
|
||||
# 9. Project outputs
|
||||
attn_output = self.attn.to_out[0](attn_output)
|
||||
context_attn_output = self.attn.to_add_out(context_attn_output)
|
||||
|
||||
# 10. Apply norm and feed forward (same as PatchedJointTransformerBlock)
|
||||
hidden_states = JointTransformerBlock.apply_norm_and_feed_forward(
|
||||
hidden_states=hidden_states,
|
||||
attn_output=attn_output,
|
||||
gate_mlp=gate_mlp,
|
||||
gate_msa=gate_msa,
|
||||
scale_mlp=scale_mlp,
|
||||
shift_mlp=shift_mlp,
|
||||
norm_layer=self.block.norm2,
|
||||
ff_layer=self.block.ff,
|
||||
)
|
||||
encoder_hidden_states = JointTransformerBlock.apply_norm_and_feed_forward(
|
||||
hidden_states=encoder_hidden_states,
|
||||
attn_output=context_attn_output,
|
||||
gate_mlp=c_gate_mlp,
|
||||
gate_msa=c_gate_msa,
|
||||
scale_mlp=c_scale_mlp,
|
||||
shift_mlp=c_shift_mlp,
|
||||
norm_layer=self.block.norm2_context,
|
||||
ff_layer=self.block.ff_context,
|
||||
)
|
||||
|
||||
return encoder_hidden_states, hidden_states
|
||||
|
||||
|
||||
class CachingSingleTransformerBlock:
|
||||
"""Single transformer block that captures K/V for cache during sync mode.
|
||||
|
||||
Implements the full block computation while storing K/V in the cache for
|
||||
subsequent async timesteps to use as stale values. Computes K/V only once.
|
||||
"""
|
||||
|
||||
def __init__(self, block: SingleTransformerBlock, kv_cache: ImagePatchKVCache):
|
||||
"""Wrap an existing SingleTransformerBlock with a KV cache.
|
||||
|
||||
Args:
|
||||
block: The original SingleTransformerBlock
|
||||
kv_cache: KV cache to populate during forward pass
|
||||
"""
|
||||
self.block = block
|
||||
self.kv_cache = kv_cache
|
||||
self.attn = block.attn
|
||||
self.num_heads = block.attn.num_heads
|
||||
self.head_dim = block.attn.head_dimension
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
text_embeddings: mx.array,
|
||||
rotary_embeddings: mx.array,
|
||||
text_seq_len: int,
|
||||
) -> mx.array:
|
||||
"""Forward pass that computes attention and populates the KV cache.
|
||||
|
||||
Args:
|
||||
hidden_states: Full [text + image] hidden states [B, text_len + img_len, D]
|
||||
text_embeddings: Time + pooled text conditioning
|
||||
rotary_embeddings: Full rotary embeddings for [text + full_image]
|
||||
text_seq_len: Length of text portion (needed to cache only image K/V)
|
||||
|
||||
Returns:
|
||||
Output hidden states after block processing
|
||||
"""
|
||||
total_seq_len = hidden_states.shape[1]
|
||||
num_img_tokens = total_seq_len - text_seq_len
|
||||
batch_size = hidden_states.shape[0]
|
||||
|
||||
# 0. Establish residual connection
|
||||
residual = hidden_states
|
||||
|
||||
# 1. Compute norm
|
||||
norm_hidden, gate = self.block.norm(
|
||||
hidden_states=hidden_states,
|
||||
text_embeddings=text_embeddings,
|
||||
)
|
||||
|
||||
# 2. Compute Q, K, V
|
||||
query, key, value = AttentionUtils.process_qkv(
|
||||
hidden_states=norm_hidden,
|
||||
to_q=self.attn.to_q,
|
||||
to_k=self.attn.to_k,
|
||||
to_v=self.attn.to_v,
|
||||
norm_q=self.attn.norm_q,
|
||||
norm_k=self.attn.norm_k,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 3. Apply RoPE to Q and K
|
||||
query, key = AttentionUtils.apply_rope(
|
||||
xq=query, xk=key, freqs_cis=rotary_embeddings
|
||||
)
|
||||
|
||||
# 4. Store only IMAGE K/V in cache for async pipeline warmstart
|
||||
# Text K/V is not cached - it's always computed fresh
|
||||
self.kv_cache.update_image_patch(
|
||||
patch_start=0,
|
||||
patch_end=num_img_tokens,
|
||||
key=key[:, :, text_seq_len:, :],
|
||||
value=value[:, :, text_seq_len:, :],
|
||||
)
|
||||
|
||||
# 5. Compute full attention (standard, not using stale cache)
|
||||
attn_output = AttentionUtils.compute_attention(
|
||||
query=query,
|
||||
key=key,
|
||||
value=value,
|
||||
batch_size=batch_size,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
|
||||
# 6. Apply feed forward and projection (same as PatchedSingleTransformerBlock)
|
||||
hidden_states = self.block._apply_feed_forward_and_projection(
|
||||
norm_hidden_states=norm_hidden,
|
||||
attn_output=attn_output,
|
||||
gate=gate,
|
||||
)
|
||||
|
||||
return residual + hidden_states
|
||||
@@ -1,7 +1,7 @@
|
||||
from exo.shared.types.worker.instances import BoundInstance
|
||||
from exo.worker.engines.mflux.distributed_flux import DistributedFlux1
|
||||
from exo.worker.engines.mflux.distributed_model import DistributedImageModel
|
||||
|
||||
|
||||
def initialize_mflux(bound_instance: BoundInstance) -> DistributedFlux1:
|
||||
"""Initialize DistributedFlux1 from a BoundInstance."""
|
||||
return DistributedFlux1.from_bound_instance(bound_instance)
|
||||
def initialize_mflux(bound_instance: BoundInstance) -> DistributedImageModel:
|
||||
"""Initialize DistributedImageModel from a BoundInstance."""
|
||||
return DistributedImageModel.from_bound_instance(bound_instance)
|
||||
|
||||
Reference in New Issue
Block a user