dataclass ¶
set forward context for diffusion models
class-attribute instance-attribute ¶
attn_metadata: (
dict[str, AttentionMetadata]
| list[dict[str, AttentionMetadata]]
| None
) = None
class-attribute instance-attribute ¶
in_diffusion_kv_memory_profile: bool = False
class-attribute instance-attribute ¶
omni_diffusion_config: OmniDiffusionConfig | None = None
class-attribute instance-attribute ¶
paged_kv_runtime: DiffusionPagedAttentionRuntime | None = (
None
)
property ¶
sp_active: bool
Returns True when SP attention parallelism should be enabled.
- If _sp_plan hooks are applied: use _sp_shard_depth (0 = outside sharded region).
- If _sp_plan hooks are NOT applied: default to True when sequence_parallel_size > 1, since _sp_shard_depth is only meaningful within the _sp_plan hook mechanism.
property ¶
sp_rank_local_seq_lens_equal: bool
Whether every active SP boundary guarantees equal local shard sizes.
Only framework-managed auto_pad boundaries provide this contract; a region that also shards manually keeps the dynamic length exchange.
class-attribute instance-attribute ¶
Build a per-rank SP padding mask that matches the local shard shape.
Auto-padding is applied before sequence-parallel sharding, so attention on each rank must receive a mask for its local shard, not for the global padded sequence.
create_forward_context(
vllm_config: VllmConfig | None = None,
omni_diffusion_config: OmniDiffusionConfig
| None = None,
attn_metadata: dict[str, AttentionMetadata]
| list[dict[str, AttentionMetadata]]
| None = None,
paged_kv_runtime: DiffusionPagedAttentionRuntime
| None = None,
paged_kv_cached_prefix_len: int = 0,
in_diffusion_kv_memory_profile: bool = False,
split_text_embed_in_sp: bool = False,
denoise_step_idx: int | None = None,
)
Pre-padding global length of shard_group, or None if it was not split.
Resolve the Ulysses-SP mode from the current ForwardContext.
Returns default when ForwardContext is unavailable or the diffusion config is not set.
override_forward_context(
forward_context: ForwardContext | None,
)
A context manager that overrides the current forward context. This is used to override the forward context for a specific forward pass.
override_paged_kv_adapter(adapter: Any | None)
Temporarily expose a Worker paged-KV adapter to Omni Attention.
This is deliberately a small context override instead of a second global forward context. The model runner owns the outer context and the adapter only replaces one opaque field while its prepared native metadata is live.
set_forward_context(
vllm_config: VllmConfig | None = None,
omni_diffusion_config: OmniDiffusionConfig
| None = None,
attn_metadata: dict[str, AttentionMetadata]
| list[dict[str, AttentionMetadata]]
| None = None,
paged_kv_runtime: DiffusionPagedAttentionRuntime
| None = None,
paged_kv_cached_prefix_len: int = 0,
in_diffusion_kv_memory_profile: bool = False,
split_text_embed_in_sp: bool = False,
denoise_step_idx: int | None = None,
)
A context manager that stores the current forward context, can be attention metadata, split_text_embed_in_sp, etc. Here we can inject common logic for every model forward pass.
set_forward_context_denoise_step_idx(
step_idx: int | None,
) -> None
Set the current diffusion denoise step on the active ForwardContext.
set_forward_context_denoise_timestep(
timestep: float | None,
) -> None
Set the normalized (descending, 1 -> 0) denoise timestep.
Timestep-gated attention features read this; pipelines that drive their own denoise loop can publish it directly instead of going through :meth:DenoiseProgressMixin.record_denoise_step.
set_forward_context_denoise_total_steps(
total_steps: int | None,
) -> None
Set the total denoise step count on the active ForwardContext.
Denoise loops publish it so tail-fallback gates (e.g. end_step in RAINFUSION_ATTN) know when the final denoise steps begin.
Set the projected direct-VLM condition on the active context.
Set the per-request reference latent on the active ForwardContext.
Used by img2img-capable DiT models (e.g. Ming-flash-omni-2.0) so the transformer can read the reference latent from request scope instead of module instance state.