From 636d3b743e9ff199dab2c8ca9eda39d022bbaca7 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Sun, 10 May 2026 11:51:10 -0700 Subject: [PATCH] [misc] attention hot-path cleanup + denoising loop hoists (#1272) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- fastvideo/attention/backends/bsa_attn.py | 16 +++-- fastvideo/attention/backends/flash_attn.py | 42 ++++++------ fastvideo/attention/backends/sla.py | 16 +++-- .../attention/backends/video_sparse_attn.py | 66 ++++++++++++++----- fastvideo/pipelines/stages/conditioning.py | 33 ++-------- fastvideo/pipelines/stages/denoising.py | 28 +++++--- 6 files changed, 115 insertions(+), 86 deletions(-) diff --git a/fastvideo/attention/backends/bsa_attn.py b/fastvideo/attention/backends/bsa_attn.py index be804f3e5..ddacdb303 100644 --- a/fastvideo/attention/backends/bsa_attn.py +++ b/fastvideo/attention/backends/bsa_attn.py @@ -358,12 +358,14 @@ def _flash_attn_single_mask( flat_k = torch.cat(k_list, dim=0) flat_v = torch.cat(v_list, dim=0) + # Compute max_seqlen_k from the Python list before moving to GPU to + # avoid a `.item()` round-trip that would force a host/device sync. + max_seqlen_q = Sq + max_seqlen_k = max(b - a for a, b in zip(cu_seqlens_k[:-1], cu_seqlens_k[1:], strict=False)) + cu_seqlens_q_t = torch.tensor(cu_seqlens_q, dtype=torch.int32, device=device) cu_seqlens_k_t = torch.tensor(cu_seqlens_k, dtype=torch.int32, device=device) - max_seqlen_q = Sq - max_seqlen_k = int((cu_seqlens_k_t[1:] - cu_seqlens_k_t[:-1]).max().item()) - orig_dtype = flat_q.dtype compute_dtype = orig_dtype if compute_dtype not in (torch.float16, torch.bfloat16): @@ -445,12 +447,14 @@ def _flash_attn_single_head( flat_k = torch.cat(k_list, dim=0) flat_v = torch.cat(v_list, dim=0) + # Compute max_seqlen_k from the Python list before moving to GPU to + # avoid a `.item()` round-trip that would force a host/device sync. + max_seqlen_q = Sq + max_seqlen_k = max(b - a for a, b in zip(cu_seqlens_k[:-1], cu_seqlens_k[1:], strict=False)) + cu_seqlens_q_t = torch.tensor(cu_seqlens_q, dtype=torch.int32, device=device) cu_seqlens_k_t = torch.tensor(cu_seqlens_k, dtype=torch.int32, device=device) - max_seqlen_q = Sq - max_seqlen_k = int((cu_seqlens_k_t[1:] - cu_seqlens_k_t[:-1]).max().item()) - orig_dtype = flat_q.dtype compute_dtype = orig_dtype if compute_dtype not in (torch.float16, torch.bfloat16): diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 80f3571b3..3ca5c047b 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -57,6 +57,24 @@ class FlashAttentionBackend(AttentionBackend): raise NotImplementedError +def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor: + # Normalize attn_mask to [B, key_len] where True means valid token. + if attn_mask.dim() == 4: + attn_mask = attn_mask[:, 0, 0, :] + elif attn_mask.dim() == 3: + attn_mask = attn_mask[:, 0, :] + elif attn_mask.dim() != 2: + raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}") + + # SDPA additive mask convention: valid=0, masked=-inf/large negative. + key_padding_mask = attn_mask if attn_mask.dtype == torch.bool else attn_mask >= 0 + + if key_padding_mask.shape[-1] != key_len: + raise ValueError("Invalid key padding mask length for FLASH_ATTN: " + f"expected {key_len}, got {key_padding_mask.shape[-1]}") + return key_padding_mask + + @dataclass class FlashAttnMetadata(AttentionMetadata): current_timestep: int @@ -101,24 +119,6 @@ class FlashAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: FlashAttnMetadata, ): - - def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor: - # Normalize attn_mask to [B, key_len] where True means valid token. - if attn_mask.dim() == 4: - attn_mask = attn_mask[:, 0, 0, :] - elif attn_mask.dim() == 3: - attn_mask = attn_mask[:, 0, :] - elif attn_mask.dim() != 2: - raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}") - - # SDPA additive mask convention: valid=0, masked=-inf/large negative. - key_padding_mask = attn_mask if attn_mask.dtype == torch.bool else attn_mask >= 0 - - if key_padding_mask.shape[-1] != key_len: - raise ValueError("Invalid key padding mask length for FLASH_ATTN: " - f"expected {key_len}, got {key_padding_mask.shape[-1]}") - return key_padding_mask - if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None): from fastvideo.attention.utils.flash_attn_no_pad import ( flash_attn_no_pad, @@ -136,6 +136,7 @@ class FlashAttentionImpl(AttentionImpl): device=query.device, ) key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, key.shape[1]).to(device=key.device) + return flash_attn_varlen_qk_no_pad( query, key, @@ -148,9 +149,8 @@ class FlashAttentionImpl(AttentionImpl): ) qkv = torch.stack([query, key, value], dim=2) - - attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True) - output = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0, softmax_scale=None) + attn_mask_padded = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True) + output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None) else: output = flash_attn_func( query, # type: ignore[no-untyped-call] diff --git a/fastvideo/attention/backends/sla.py b/fastvideo/attention/backends/sla.py index 59b232b3a..efb3cf882 100644 --- a/fastvideo/attention/backends/sla.py +++ b/fastvideo/attention/backends/sla.py @@ -307,9 +307,12 @@ class SLAAttentionImpl(AttentionImpl, nn.Module): # Sparse attention o_s = _attention.apply(q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK) - # Linear attention with feature maps - q_linear = self.feature_map_q(q).contiguous().to(self.dtype) - k_linear = self.feature_map_k(k).contiguous().to(self.dtype) + # Linear attention with feature maps. Note: softmax / elu / relu + # are elementwise and preserve layout, so the inputs are already + # contiguous from the transpose-contiguous above — no need to + # call .contiguous() again here. + q_linear = self.feature_map_q(q).to(self.dtype) + k_linear = self.feature_map_k(k).to(self.dtype) o_l = self._calc_linear_attention(q_linear, k_linear, v) # Project linear attention output and combine @@ -539,9 +542,10 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module): False, 1, scale, 0) # ========== END SPARGE ========== - # Linear attention with feature maps - q_linear = self.feature_map_q(q).contiguous().to(self.dtype) - k_linear = self.feature_map_k(k).contiguous().to(self.dtype) + # Linear attention with feature maps (see SLAAttentionImpl.forward + # for why .contiguous() is unnecessary here). + q_linear = self.feature_map_q(q).to(self.dtype) + k_linear = self.feature_map_k(k).to(self.dtype) o_l = self._calc_linear_attention(q_linear, k_linear, v) # Project linear attention output and combine diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index 43960a940..3ab483222 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -140,6 +140,18 @@ class VideoSparseAttentionMetadata(AttentionMetadata): reverse_tile_partition_indices: torch.LongTensor variable_block_sizes: torch.LongTensor non_pad_index: torch.LongTensor + # Precomputed fancy index that fuses ``x[:, non_pad_index][:, reverse_tile_partition_indices]`` + # in postprocess_output(). Avoids materializing the intermediate + # ``[B, len(non_pad_index), H, D]`` tensor on every layer. + untile_combined_index: torch.LongTensor + # Per-step shared padded buffer used by tile(). Lazily populated on + # the first layer's call and reused by every subsequent VSA layer in + # the same denoising step. Scoping to metadata (not class/instance) + # makes the reuse thread-safe across concurrent requests and keeps + # the "pad positions are zero" invariant trivially true (the buffer + # is freshly zeroed alongside ``non_pad_index`` so the index set + # cannot drift between calls). + tile_buf: torch.Tensor | None = None class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): @@ -171,6 +183,7 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): reverse_tile_partition_indices = get_reverse_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device) variable_block_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device) non_pad_index = get_non_pad_index(variable_block_sizes, math.prod(VSA_TILE_SIZE)) + untile_combined_index = non_pad_index[reverse_tile_partition_indices] return VideoSparseAttentionMetadata( current_timestep=current_timestep, @@ -181,7 +194,8 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): tile_partition_indices=tile_partition_indices, # type: ignore reverse_tile_partition_indices=reverse_tile_partition_indices, variable_block_sizes=variable_block_sizes, - non_pad_index=non_pad_index) + non_pad_index=non_pad_index, + untile_combined_index=untile_combined_index) class VideoSparseAttentionImpl(AttentionImpl): @@ -200,37 +214,59 @@ class VideoSparseAttentionImpl(AttentionImpl): sp_group = get_sp_group() self.sp_size = sp_group.world_size - def tile(self, x: torch.Tensor, num_tiles: list[int], tile_partition_indices: torch.LongTensor, - non_pad_index: torch.LongTensor) -> torch.Tensor: + def tile(self, x: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata) -> torch.Tensor: + """Tile ``x`` into ``attn_metadata.tile_buf`` and return it. + + The returned tensor aliases the per-metadata buffer and is only + valid until the next ``tile()`` / ``preprocess_qkv`` call on the + same ``attn_metadata``. Callers must consume (or copy) the + result before invoking another VSA layer with the same metadata. + Today both call sites materialize copies via + ``.transpose(...).contiguous()`` inside ``forward()``, so the + contract holds; future callers must preserve it. + """ + num_tiles = attn_metadata.num_tiles t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0] h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1] w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2] + target_shape = (x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1]) - x_padded = torch.zeros((x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1]), - device=x.device, - dtype=x.dtype) - x_padded[:, non_pad_index] = x[:, tile_partition_indices] - return x_padded + # Reuse the per-step buffer stashed on metadata (lazily allocated + # on the first VSA layer's call within a denoising step). Pad + # positions are zero from the initial torch.zeros and never + # written to. Scoping to metadata makes reuse safe across + # concurrent requests and keeps the "pad positions are zero" + # invariant trivially true: ``non_pad_index`` is fixed within + # a single metadata instance. + buf = attn_metadata.tile_buf + if (buf is None or buf.shape != target_shape or buf.dtype != x.dtype or buf.device != x.device): + buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype) + attn_metadata.tile_buf = buf - def untile(self, x: torch.Tensor, reverse_tile_partition_indices: torch.LongTensor, - non_pad_index: torch.LongTensor) -> torch.Tensor: - x = x[:, non_pad_index][:, reverse_tile_partition_indices] - return x + buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices] + return buf + + def untile(self, x: torch.Tensor, untile_combined_index: torch.LongTensor) -> torch.Tensor: + # Single fancy index using precomputed combined indices; avoids + # the intermediate ``[B, len(non_pad_index), H, D]`` tensor that + # the two-step ``x[:, non_pad_index][:, reverse_tile_partition_indices]`` + # would allocate on every layer. + return x[:, untile_combined_index] def preprocess_qkv( self, qkv: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: - return self.tile(qkv, attn_metadata.num_tiles, attn_metadata.tile_partition_indices, - attn_metadata.non_pad_index) + """Tile QKV; aliasing contract: see ``tile()``.""" + return self.tile(qkv, attn_metadata) def postprocess_output( self, output: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: - return self.untile(output, attn_metadata.reverse_tile_partition_indices, attn_metadata.non_pad_index) + return self.untile(output, attn_metadata.untile_combined_index) def forward( # type: ignore[override] self, diff --git a/fastvideo/pipelines/stages/conditioning.py b/fastvideo/pipelines/stages/conditioning.py index 80fb3b6cc..95ee9c0d4 100644 --- a/fastvideo/pipelines/stages/conditioning.py +++ b/fastvideo/pipelines/stages/conditioning.py @@ -6,14 +6,11 @@ Conditioning stage for diffusion pipelines. import torch from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult -logger = init_logger(__name__) - class ConditioningStage(PipelineStage): """ @@ -39,31 +36,11 @@ class ConditioningStage(PipelineStage): Returns: The batch with applied conditioning. """ - # TODO!! - if not batch.do_classifier_free_guidance: - return batch - else: - return batch - - logger.info("batch.negative_prompt_embeds: %s", batch.negative_prompt_embeds) - logger.info("do_classifier_free_guidance: %s", batch.do_classifier_free_guidance) - logger.info("cfg_scale: %s", batch.guidance_scale) - - # Ensure negative prompt embeddings are available - assert batch.negative_prompt_embeds is not None, ( - "Negative prompt embeddings are required for classifier-free guidance") - - # Concatenate primary embeddings and masks - batch.prompt_embeds = torch.cat([batch.negative_prompt_embeds, batch.prompt_embeds]) - if batch.attention_mask is not None: - batch.attention_mask = torch.cat([batch.negative_attention_mask, batch.attention_mask]) - - # Concatenate secondary embeddings and masks if present - if batch.prompt_embeds_2 is not None: - batch.prompt_embeds_2 = torch.cat([batch.negative_prompt_embeds_2, batch.prompt_embeds_2]) - if batch.attention_mask_2 is not None: - batch.attention_mask_2 = torch.cat([batch.negative_attention_mask_2, batch.attention_mask_2]) - + # Forward is a no-op: CFG is applied via two separate + # transformer forward passes inside DenoisingStage (e.g. + # denoising.py:364-394, :706, :930). The class is kept because + # verify_input / verify_output still validate CFG fields and + # disable CFG when prompt_embeds is empty. return batch def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index c4244e9ac..c112f48ca 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -212,6 +212,23 @@ class DenoisingStage(PipelineStage): trajectory_timesteps: list[torch.Tensor] = [] trajectory_latents: list[torch.Tensor] = [] + # Hoisted out of the per-step loop: depends only on inputs that + # are constant across denoising steps. + use_meanflow = getattr(self.transformer.config, "use_meanflow", False) + embedded_cfg_scale = fastvideo_args.pipeline_config.embedded_cfg_scale + if embedded_cfg_scale is not None: + guidance_expand = (torch.tensor( + [embedded_cfg_scale] * latents.shape[0], + dtype=torch.float32, + device=get_local_torch_device(), + ).to(target_dtype) * 1000.0) + else: + guidance_expand = None + # V2V padding: zero-filled tensor concatenated with each step's + # latent_model_input. Shape is fixed by latents and is never + # written to, so we allocate once. + v2v_zero_pad = torch.zeros_like(latents) if batch.video_latent is not None else None + # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -248,8 +265,7 @@ class DenoisingStage(PipelineStage): # Expand latents for V2V/I2V latent_model_input = latents.to(target_dtype) if batch.video_latent is not None: - latent_model_input = torch.cat([latent_model_input, batch.video_latent, - torch.zeros_like(latents)], + latent_model_input = torch.cat([latent_model_input, batch.video_latent, v2v_zero_pad], dim=1).to(target_dtype) elif batch.image_latent is not None: assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task" @@ -266,7 +282,6 @@ class DenoisingStage(PipelineStage): t_expand = t.repeat(latent_model_input.shape[0]) t_expand = t_expand.to(get_local_torch_device()) - use_meanflow = getattr(self.transformer.config, "use_meanflow", False) if use_meanflow: if i == len(timesteps) - 1: timesteps_r = torch.tensor([0.0], device=get_local_torch_device()) @@ -285,13 +300,6 @@ class DenoisingStage(PipelineStage): latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - # Prepare inputs for transformer - guidance_expand = (torch.tensor( - [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0], - dtype=torch.float32, - device=get_local_torch_device(), - ).to(target_dtype) * 1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale is not None else None) - # Predict noise residual with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): if (vsa_available and self.attn_backend == VideoSparseAttentionBackend):