From 5ab2cfe9a899d1d2dffae932fb0f6371014d0d71 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 14 Mar 2026 23:41:23 +0800 Subject: [PATCH] [Diffusion] Clean upstream fa3 in hopper (#20576) --- .../scripts/bench_diffusion_denoise.py | 2 +- .../scripts/bench_diffusion_rmsnorm.py | 2 +- .../layers/attention/backends/flash_attn.py | 106 ------------------ 3 files changed, 2 insertions(+), 108 deletions(-) diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py index 4ef1a6c46..59da0e110 100755 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py @@ -8,7 +8,7 @@ Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kerne Usage: # Baseline — single model - cd /data/bbuf/sglang + cd /path/to/sglang python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux # With custom JIT CUDA kernels diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py index b73d7bd6e..4bdadcf7c 100755 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py @@ -8,7 +8,7 @@ Compares: Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels Usage: - cd /data/bbuf/sglang + cd /path/to/sglang python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py Requirements: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index 72100c65b..9c30a9798 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -1,7 +1,6 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass -from functools import lru_cache from typing import Any, List, Optional, Tuple import torch @@ -10,7 +9,6 @@ from sglang.multimodal_gen.runtime.layers.utils import register_custom_op from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, - current_platform, ) try: @@ -28,10 +26,6 @@ try: except ImportError as e: raise e -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) - def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return x.contiguous() if x is not None and x.stride(-1) != 1 else x @@ -306,20 +300,6 @@ def flash_attn_varlen_func_op_lse( ) -try: - if current_platform.is_hopper(): - from flash_attn_interface import ( - flash_attn_varlen_func as flash_attn_varlen_func_upstream, - ) - else: - flash_attn_varlen_func_upstream = None - -except Exception: - flash_attn_varlen_func_upstream = None - logger.warning( - "flash_attn 3 package is not installed. It's recommended to install flash_attn3 on hopper, otherwise performance is sub-optimal" - ) - from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionBackend, AttentionImpl, @@ -330,51 +310,6 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i fa_ver = 3 -@lru_cache(maxsize=128) -def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor: - return torch.arange( - 0, - (bsz + 1) * seqlen, - step=seqlen, - device=torch.device("cuda", device_index), - dtype=torch.int32, - ) - - -@lru_cache(maxsize=256) -def _should_use_upstream_flash_attention( - upstream_available: bool, - upstream_heads_ok: bool, - q_shape: tuple[int, ...], - k_shape: tuple[int, ...], - v_shape: tuple[int, ...], -) -> bool: - if not upstream_available or not upstream_heads_ok: - return False - - if len(q_shape) != 4 or len(k_shape) != 4 or len(v_shape) != 4: - return False - - bsz, seqlen, nheads_q, d = q_shape - bsz_k, seqlen_k, nheads_k, d_k = k_shape - bsz_v, seqlen_v, nheads_v, d_v = v_shape - - if ( - bsz != bsz_k - or bsz != bsz_v - or seqlen != seqlen_k - or seqlen != seqlen_v - or d != d_k - or d != d_v - ): - return False - if nheads_k != nheads_v: - return False - if nheads_k == 0 or (nheads_q % nheads_k) != 0: - return False - return True - - def set_fa_ver(ver: int) -> None: global fa_ver fa_ver = ver @@ -450,13 +385,6 @@ class FlashAttentionImpl(AttentionImpl): self.causal = causal self.softmax_scale = softmax_scale self.attention_metadata = FlashAttentionMetadata() - if self.num_kv_heads is None: - self._upstream_heads_ok = True - else: - # For gqa, the num_heads must be a multiple of num_kv_heads - self._upstream_heads_ok = ( - self.num_kv_heads > 0 and (self.num_heads % self.num_kv_heads) == 0 - ) def forward( self, @@ -477,40 +405,6 @@ class FlashAttentionImpl(AttentionImpl): max_seqlen_q = query.shape[1] max_seqlen_k = key.shape[1] - q_shape = tuple(query.shape) - k_shape = tuple(key.shape) - v_shape = tuple(value.shape) - - use_upstream = _should_use_upstream_flash_attention( - flash_attn_varlen_func_upstream is not None, - self._upstream_heads_ok, - q_shape, - k_shape, - v_shape, - ) - - if use_upstream: - bsz, seqlen, nheads_q, d = q_shape - q_ = query.contiguous() - k_ = key.contiguous() - v_ = value.contiguous() - out = flash_attn_varlen_func_upstream( - q_, - k_, - v_, - None, - None, - seqlen, - seqlen, - softmax_scale=self.softmax_scale, - causal=self.causal, - return_attn_probs=return_softmax_lse, - ) - if return_softmax_lse: - out_tensor, softmax_lse = out - return out_tensor.reshape(bsz, seqlen, nheads_q, -1), softmax_lse - return out.reshape(bsz, seqlen, nheads_q, d) - # FA version selection: # - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag) # - fa_ver == 4: call custom ops with FIXED return schema