Merge pull request #3914 from vladmandic/dev

merge dev

.pylintrc

@@ -23,6 +23,7 @@ ignore-paths=/usr/lib/.*$,
              modules/k-diffusion,
              modules/flex2,
              modules/ldsr,
+             modules/hidream,
              modules/meissonic,
              modules/mod,
              modules/omnigen,

.ruff.toml

@@ -20,6 +20,7 @@ exclude = [
     "modules/meissonic",
     "modules/mod",
     "modules/omnigen",
+    "modules/hidream",
     "modules/pag",
     "modules/pixelsmith",
     "modules/postprocess/aurasr_arch.py",

CHANGELOG.md

@@ -1,5 +1,44 @@
 # Change Log for SD.Next
 
+## Update for 2025-05-06
+
+- **Features**  
+  - [FramePack](https://vladmandic.github.io/sdnext-docs/FramePack)  
+    add **T2V** mode in addition to **I2V** and **FLF2V**  
+    support for new **F1: forward-only** model variant in addition to regular **bi-directional**  
+    add **prompt enhance** using VLM: it will analyze input image and then create enhanced prompt based on user prompt and image  
+    add **prompt interpolation**, section prompts do not need to match exact video section count  
+    and improved performance  
+    [Docs](https://vladmandic.github.io/sdnext-docs/FramePack) rewrite!  
+  - **Prompt-Enhhance**  
+    add **Qwen3** *0.6B/1.7B/4B* models  
+    add thinking mode support (for models that have it)  
+  - [HiDream-E1](https://huggingface.co/HiDream-ai/HiDream-E1-Full) natural language image-editing model built on HiDream-I1  
+    available via  *networks -> models -> reference*  
+    *note*: right now HiDream-E1 is limited to 768x768 images, so you must force resize image before running it  
+- **Other**  
+  - CUDA: set default to `torch==2.7.0` with `cuda==12.8`  
+  - ZLUDA: update to `zluda==3.9.4` and `flash-attn-2`  
+  - Docker: pre-install `ffmpeg`  
+  - Wiki: updated pages: *FramePack, Video, ROCm, ZLUDA, Quantization*  
+  - Gallery: support JXL image format  
+  - Scheduler: add sigmoid beta scheduler  
+  - GitHub: updated issue template  
+- **Fixes**  
+  - FramePack: correct dtype  
+  - NNCF: check dependencies and register quant type  
+  - API: refresh checkpoint list  
+  - API: vlm-api endpoint  
+  - Styles: save style with prompt  
+  - Texture tiling: fix apply when switching models  
+  - Diffusers: slow initial startup  
+  - Gated access: obfuscate and log token used for access  
+  - SDXL refiner workflow  
+  - Control: t2i-adapter workflow  
+  - Control: xs-controlnet workflow  
+  - Control: lllite-workflow  
+  - Control: refiner workflow with multiple control elements  
+
 ## Highlights for 2025-04-28
 
 Another major release with *over 120 commits*!  

TODO.md

@@ -4,6 +4,7 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 ## Current
 
+- Video: API support  
 - ModernUI for custom model loader  
 - ModernUI for history tab  
 
@@ -13,10 +14,8 @@ N/A
 
 ## Future Candidates
 
-- IPAdapter: negative guidance: <https://github.com/huggingface/diffusers/discussions/7167>  
 - Control: API enhance scripts compatibility  
-- Video: add generate context menu  
-- Video: API support  
+- IPAdapter: negative guidance: <https://github.com/huggingface/diffusers/discussions/7167>  
 - Video: STG: <https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance>  
 - Video: SmoothCache: https://github.com/huggingface/diffusers/issues/11135  
 
@@ -27,9 +26,7 @@ N/A
 - control: support scripts via api
 - fc: autodetect distilled based on model
 - fc: autodetect tensor format based on model
-- hidream: pack latents for remote vae
 - hypertile: vae breaks when using non-standard sizes
-- infotext: handle using regex instead
 - install: enable ROCm for windows when available
 - loader: load receipe
 - loader: save receipe

configs/Dockerfile.cuda

@@ -18,7 +18,7 @@ LABEL org.opencontainers.image.version="latest"
 
 # minimum install
 RUN ["apt-get", "-y", "update"]
-RUN ["apt-get", "-y", "install", "git", "build-essential", "google-perftools", "curl"]
+RUN ["apt-get", "-y", "install", "git", "build-essential", "google-perftools", "curl", "ffmpeg"]
 # optional if full cuda-dev is required by some downstream library
 # RUN ["apt-get", "-y", "nvidia-cuda-toolkit"]
 RUN ["/usr/sbin/ldconfig"]

html/reference.json

@@ -364,6 +364,13 @@
     "skip": true,
     "extras": "sampler: Default"
   },
+  "HiDream-E1 Full": {
+    "path": "HiDream-ai/HiDream-E1-Full",
+    "desc": "HiDream-E1 is an image editing model built on HiDream-I1.",
+    "preview": "HiDream-ai--HiDream-I1-Fast.jpg",
+    "skip": true,
+    "extras": "sampler: Default"
+  },
 
   "Kwai Kolors": {
     "path": "Kwai-Kolors/Kolors-diffusers",

installer.py

@@ -508,17 +508,23 @@ def get_platform():
 
 
 # check python version
-def check_python(supported_minors=[9, 10, 11, 12], reason=None):
+def check_python(supported_minors=[], experimental_minors=[], reason=None):
+    if supported_minors is None or len(supported_minors) == 0:
+        supported_minors = [9, 10, 11, 12]
+        experimental_minors = []
     t_start = time.time()
     if args.quick:
         return
     log.info(f'Python: version={platform.python_version()} platform={platform.system()} bin="{sys.executable}" venv="{sys.prefix}"')
     if not (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in supported_minors):
-        log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
-        if reason is not None:
-            log.error(reason)
-        if not args.ignore and not args.experimental:
-            sys.exit(1)
+        if (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in experimental_minors):
+            log.warning(f"Python version experimental: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} recommended 3.{supported_minors}")
+        else:
+            log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
+            if reason is not None:
+                log.error(reason)
+            if not args.ignore and not args.experimental:
+                sys.exit(1)
     if int(sys.version_info.minor) == 12:
         os.environ.setdefault('SETUPTOOLS_USE_DISTUTILS', 'local') # hack for python 3.11 setuptools
     if not args.skip_git:
@@ -538,7 +544,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = 'b4be42282dc9bf9deccc2f60f5db952de772cf42' # diffusers commit hash
+    sha = '8c661ea586bf11cb2440da740dd3c4cf84679b85' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
     cur = opts.get('diffusers_version', '') if minor > 0 else ''
@@ -573,7 +579,8 @@ def install_cuda():
     if args.use_nightly:
         cmd = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 --extra-index-url https://download.pytorch.org/whl/nightly/cu126')
     else:
-        cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126')
+        # cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126')
+        cmd = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+cu128 torchvision==0.22.0+cu128 --index-url https://download.pytorch.org/whl/cu128')
     return cmd
 
 
@@ -1090,7 +1097,13 @@ def install_submodules(force=True):
     return '\n'.join(res)
 
 
-def reload(package):
+def reload(package, desired=None):
+    loaded = package in sys.modules
+    if not loaded:
+        return
+    current = sys.modules[package].__version__ if hasattr(sys.modules[package], "__version__") else None
+    if desired is not None and current == desired:
+        return
     modules = [m for m in sys.modules if m.startswith(package)]
     for m in modules:
         del sys.modules[m]
@@ -1142,17 +1155,17 @@ def install_optional():
     install('pillow-jxl-plugin==1.3.2', ignore=True)
     install('optimum-quanto==0.2.7', ignore=True)
     install('torchao==0.10.0', ignore=True)
-    install('bitsandbytes==0.45.1', ignore=True)
+    install('bitsandbytes==0.45.5', ignore=True)
     install('pynvml', ignore=True)
     install('ultralytics==8.3.40', ignore=True)
     install('Cython', ignore=True)
     install('insightface==0.7.3', ignore=True) # problematic build
     install('albumentations==1.4.3', ignore=True)
     install('pydantic==1.10.21', ignore=True)
-    reload('pydantic')
-    install('nncf==2.16.0', ignore=True, no_deps=True) # requires older pandas
-    # install('flash-attn', ignore=True) # requires cuda and nvcc to be installed
+    reload('pydantic', '1.10.21')
+    install('nncf==2.16.0', ignore=True) # requires older pandas
     install('gguf', ignore=True)
+    install('av', ignore=True)
     try:
         import gguf
         scripts_dir = os.path.join(os.path.dirname(gguf.__file__), '..', 'scripts')

javascript/gallery.js

@@ -126,7 +126,7 @@ class GalleryFile extends HTMLElement {
   async connectedCallback() {
     if (this.shadow.children.length > 0) return;
     const ext = this.name.split('.').pop().toLowerCase();
-    if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'svg', 'mp4'].includes(ext)) return;
+    if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) return;
     this.hash = await getHash(`${this.folder}/${this.name}/${this.size}/${this.mtime}`); // eslint-disable-line no-use-before-define
     const style = document.createElement('style');
     const width = opts.browser_fixed_width ? `${opts.extra_networks_card_size}px` : 'unset';

modules/api/endpoints.py

@@ -17,8 +17,11 @@ def get_upscalers():
     return [{"name": upscaler.name, "model_name": upscaler.scaler.model_name, "model_path": upscaler.data_path, "model_url": None, "scale": upscaler.scale} for upscaler in shared.sd_upscalers]
 
 def get_sd_models():
-    from modules import sd_models, sd_models_config
-    return [{"title": x.title, "model_name": x.name, "filename": x.filename, "type": x.type, "hash": x.shorthash, "sha256": x.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(x)} for x in sd_models.checkpoints_list.values()]
+    from modules import sd_checkpoint, sd_models_config
+    checkpoints = []
+    for v in sd_checkpoint.checkpoints_list.values():
+        checkpoints.append({"title": v.title, "model_name": v.name, "filename": v.filename, "type": v.type, "hash": v.shorthash, "sha256": v.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(v)})
+    return checkpoints
 
 def get_hypernetworks():
     return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
@@ -128,10 +131,12 @@ def post_reload_checkpoint():
     return {}
 
 def post_refresh_checkpoints():
-    return shared.refresh_checkpoints()
+    shared.refresh_checkpoints()
+    return {}
 
 def post_refresh_vae():
-    return shared.refresh_vaes()
+    shared.refresh_vaes()
+    return {}
 
 def post_refresh_loras():
     from modules.lora import lora_load

modules/api/helpers.py

@@ -24,8 +24,8 @@ def decode_base64_to_image(encoding, quiet=False):
         return image
     except Exception as e:
         shared.log.warning(f'API cannot decode image: {e}')
-        from modules import errors
-        errors.display(e, 'API cannot decode image')
+        # from modules import errors
+        # errors.display(e, 'API cannot decode image')
         if not quiet:
             raise HTTPException(status_code=500, detail="Invalid encoded image") from e
         return None

modules/api/models.py

@@ -355,6 +355,7 @@ class ReqVQA(BaseModel):
     image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
     model: str = Field(default="Microsoft Florence 2 Base", title="Model", description="The interrogate model used.")
     question: str = Field(default="describe the image", title="Question", description="Question to ask the model.")
+    system: str = Field(default="You are image captioning expert, creative, unbiased and uncensored.", title="System prompt", description="Prompt to shape how the model interprets and responds to user prompts.")
 
 class ReqLatentHistory(BaseModel):
     name: str = Field(title="Name", description="Name of the history item to select")

modules/cmd_args.py

@@ -39,6 +39,7 @@ def main_args():
     group_diag.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
     group_diag.add_argument("--monitor", default=os.environ.get("SD_MONITOR", 0), help="Run memory monitor, default: %(default)s")
     group_diag.add_argument("--status", default=os.environ.get("SD_STATUS", 120), help="Run server is-alive status, default: %(default)s")
+    group_diag.add_argument('--experimental', default=os.environ.get("SD_EXPERIMENTAL",False), action='store_true', help="Allow unsupported versions of libraries, default: %(default)s")
 
     group_http = parser.add_argument_group('HTTP')
     group_http.add_argument('--theme', type=str, default=os.environ.get("SD_THEME", None), help='Override UI theme')

modules/control/run.py

@@ -133,7 +133,7 @@ def check_active(p, unit_type, units):
         if u.type != unit_type:
             continue
         num_units += 1
-        debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled}')
+        debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled} cn={u.controlnet} proc={u.process}')
         if not u.enabled:
             if u.controlnet is not None and u.controlnet.model is not None:
                 debug_log(f'Control unit offload: model="{u.controlnet.model_id}" device={devices.cpu}')
@@ -692,7 +692,16 @@ def control_run(state: str = '',
                                 p.task_args['image'] = p.init_images # need to set explicitly for txt2img
                                 del p.init_images
                         if unit_type == 'lite':
-                            p.init_images = [input_image]
+                            if input_type == 0:
+                                shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
+                                shared.sd_model.no_task_switch = True
+                            elif input_type == 1:
+                                p.init_images = [input_image]
+                            elif input_type == 2:
+                                if init_image is None:
+                                    shared.log.warning('Control: separate init image not provided')
+                                    init_image = input_image
+                                p.init_images = [init_image]
                             instance.apply(selected_models, processed_image, control_conditioning)
                     if hasattr(p, 'init_images') and p.init_images is None: # delete empty
                         del p.init_images

modules/control/units/lite.py

@@ -125,6 +125,7 @@ class ControlLLLite():
 class ControlLLitePipeline():
     def __init__(self, pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline]):
         self.pipeline = pipeline
+        # self.pipeline.__class__.__name__ = 'ControlLLLitePipeline'
         self.nets = []
 
     def apply(self, controlnet: Union[ControlNetLLLite, list[ControlNetLLLite]], image, conditioning):

modules/control/units/t2iadapter.py

@@ -2,7 +2,7 @@ import os
 import time
 from typing import Union
 import threading
-from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, MultiAdapter, StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline # pylint: disable=unused-import
+from diffusers import pipelines, StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, MultiAdapter, StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline # pylint: disable=unused-import
 from installer import log
 from modules import errors, sd_models
 from modules.control.units import detect
@@ -142,6 +142,14 @@ class AdapterPipeline():
         if isinstance(adapter, list) and len(adapter) > 1:
             adapter = MultiAdapter(adapter)
         adapter.to(device=pipeline.device, dtype=pipeline.dtype)
+        """
+        pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["sd-t2iadapter"] = StableDiffusionAdapterPipeline
+        pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["sd-t2iadapter"] = StableDiffusionAdapterPipeline
+        pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["sd-t2iadapter"] = StableDiffusionAdapterPipeline
+        pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["sdxl-t2iadapter"] = StableDiffusionXLAdapterPipeline
+        pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["sdxl-t2iadapter"] = StableDiffusionXLAdapterPipeline
+        pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["sdxl-t2iadapter"] = StableDiffusionXLAdapterPipeline
+        """
         if pipeline.__class__.__name__ == 'StableDiffusionAdapterPipeline' or pipeline.__class__.__name__ == 'StableDiffusionXLAdapterPipeline':
             pass # already initialized
         if detect.is_sdxl(pipeline):

modules/control/units/xs_model.py

@@ -631,6 +631,7 @@ class ControlNetXSModel(ModelMixin, ConfigMixin):
         t_emb = t_emb.to(dtype=sample.dtype)
 
         if self.config.learn_embedding: # pylint: disable=no-member
+            base_model = base_model.to(self.control_model.device)
             ctrl_temb = self.control_model.time_embedding(t_emb, timestep_cond)
             base_temb = base_model.time_embedding(t_emb, timestep_cond)
             interpolation_param = self.config.time_embedding_mix**0.3 # pylint: disable=no-member

modules/control/units/xs_pipe.py

@@ -114,6 +114,14 @@ class StableDiffusionXLControlNetXSPipeline(
     ):
         super().__init__()
 
+        if isinstance(controlnet, list):
+            if len(controlnet) == 1:
+                controlnet = controlnet[0]
+            else:
+                raise ValueError(
+                    "ControlNetXS pipeline only supports a single ControlNetXS model"
+                )
+
         vae_compatible, cnxs_condition_downsample_factor, vae_downsample_factor = controlnet._check_if_vae_compatible(
             vae
         )

modules/flash_attn_triton_amd/bwd_prefill.py

@@ -1,606 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-from modules.flash_attn_triton_amd.utils import get_shape_from_layout, get_strides_from_layout
-
-
-@triton.jit
-def _bwd_preprocess_use_o(
-    Out,
-    DO,
-    Delta,
-    stride_oz, stride_oh, stride_om, stride_ok,
-    stride_doz, stride_doh, stride_dom, stride_dok, # pylint: disable=unused-argument
-    stride_deltaz, stride_deltah, stride_deltam,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    max_seqlen_q,
-    max_seqlen_k,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    ACTUAL_BLOCK_DMODEL: tl.constexpr,
-    N_CTX_Q: tl.constexpr,
-    Z: tl.constexpr, # pylint: disable=unused-argument
-    H: tl.constexpr,
-    IS_VARLEN: tl.constexpr
-):
-    pid_m = tl.program_id(0)
-    pid_bh = tl.program_id(1)
-
-    # Compute batch and head indices
-    off_z = pid_bh // H
-    off_h = pid_bh % H
-
-    if IS_VARLEN:
-        # Compute sequence lengths for the current batch
-        q_start = tl.load(cu_seqlens_q + off_z)
-        q_end = tl.load(cu_seqlens_q + off_z + 1)
-        k_start = tl.load(cu_seqlens_k + off_z)
-        k_end = tl.load(cu_seqlens_k + off_z + 1)
-
-        # Compute actual sequence lengths
-        N_CTX_Q = q_end - q_start
-        N_CTX_K = k_end - k_start # pylint: disable=unused-variable
-    else:
-        q_start = 0
-        k_start = 0
-        N_CTX_Q = max_seqlen_q
-        N_CTX_K = max_seqlen_k # pylint: disable=unused-variable
-
-    off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    off_d = tl.arange(0, BLOCK_DMODEL)
-
-    # create masks
-    mask_m = off_m < N_CTX_Q
-    mask_d = off_d < ACTUAL_BLOCK_DMODEL
-
-    # compute offsets
-    o_offset = Out + off_z * stride_oz + off_h * stride_oh + q_start * stride_om
-    do_offset = DO + off_z * stride_oz + off_h * stride_oh + q_start * stride_om
-
-    # compute pointers
-    out_ptrs = o_offset + off_m[:, None] * stride_om + off_d[None, :] * stride_ok
-    do_ptrs = do_offset + off_m[:, None] * stride_dom + off_d[None, :] * stride_dok
-
-    # load
-    o = tl.load(out_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0).to(tl.float32)
-    do = tl.load(do_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0).to(tl.float32)
-
-    # compute delta
-    delta = tl.sum(o * do, axis=1)
-
-    # write-back delta
-    delta_offset = Delta + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam
-    delta_ptrs = delta_offset + off_m * stride_deltam
-    tl.store(delta_ptrs, delta, mask=mask_m)
-
-
-@triton.jit
-def _bwd_kernel_one_col_block(
-    Q,
-    K,
-    V,
-    sm_scale,
-    Out, DO, DQ, DK, DV, L, D, # pylint: disable=unused-argument
-    q_offset,
-    k_offset,
-    v_offset,
-    do_offset,
-    dq_offset,
-    dk_offset,
-    dv_offset,
-    d_offset,
-    l_offset,
-    stride_dq_all, stride_qz, stride_qh, # pylint: disable=unused-argument
-    stride_qm,
-    stride_qk,
-    stride_kz, stride_kh, # pylint: disable=unused-argument
-    stride_kn,
-    stride_kk,
-    stride_vz, stride_vh, # pylint: disable=unused-argument
-    stride_vn,
-    stride_vk,
-    stride_deltaz,  stride_deltah, # pylint: disable=unused-argument
-    stride_deltam,
-    Z, H, # pylint: disable=unused-argument
-    N_CTX_Q,
-    N_CTX_K,
-    off_h, off_z, off_hz, # pylint: disable=unused-argument
-    start_n,
-    num_block_m,
-    num_block_n, # pylint: disable=unused-argument
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    ACTUAL_BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    SEQUENCE_PARALLEL: tl.constexpr,
-    CAUSAL: tl.constexpr,
-    USE_EXP2: tl.constexpr,
-):
-    if CAUSAL:
-        # TODO: Causal can skip more blocks with something like lo = start_m * BLOCK_M
-        lo = 0
-    else:
-        lo = 0
-
-    # initialize col and head offsets
-    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_DMODEL)
-
-    # masks
-    mask_n = offs_n < N_CTX_K
-    mask_d = offs_d < ACTUAL_BLOCK_DMODEL
-    kv_mask = mask_n[:, None] & mask_d[None, :]
-
-    # initialize grad accumulators
-    dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
-    dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
-
-    # load k and v once per column block
-    k_ptrs = k_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk
-    v_ptrs = v_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vk
-    k = tl.load(k_ptrs, mask=kv_mask, other=0.0)
-    v = tl.load(v_ptrs, mask=kv_mask, other=0.0)
-
-    # loop over rows
-    for start_m in range(lo, num_block_m * BLOCK_M, BLOCK_M):
-        offs_m = start_m + tl.arange(0, BLOCK_M)
-        q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
-        dq_ptrs = dq_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
-        do_ptrs = do_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
-
-        # update mask as row block changes
-        mask_m = offs_m < N_CTX_Q
-        q_mask = mask_m[:, None] & mask_d[None, :]
-
-        # load q, k, v, do on-chip
-        q = tl.load(q_ptrs, mask=q_mask, other=0.0)
-        do = tl.load(do_ptrs, mask=q_mask, other=0.0)
-
-        # recompute p = softmax(qk, dim=-1).T
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, tl.trans(k))
-
-        if CAUSAL:
-            col_offset = N_CTX_Q - N_CTX_K
-            causal_mask = offs_m[:, None] >= (col_offset + offs_n[None, :])
-            qk = tl.where(causal_mask, qk, float("-inf"))
-
-        l_ptrs = l_offset + offs_m * stride_deltam
-        l_i = tl.load(l_ptrs, mask=mask_m)
-
-        # compute p
-        if USE_EXP2:
-            RCP_LN2: tl.constexpr = 1.4426950408889634
-            qk *= sm_scale * RCP_LN2
-            l_i *= RCP_LN2
-            p = tl.math.exp2(qk - l_i[:, None])
-        else:
-            qk *= sm_scale
-            p = tl.math.exp(qk - l_i[:, None])
-
-        # mask block in the cases where the data is smaller the block size
-        p_mask = mask_m[:, None] & mask_n[None, :]
-        p = tl.where(p_mask, p, 0.0)
-
-        # compute dv
-        dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
-
-        # compute dp
-        dp = tl.dot(do, tl.trans(v))
-
-        # compute ds , ds = p * (dp - delta[:, None])
-        d_ptrs = d_offset + offs_m * stride_deltam
-        Di = tl.load(d_ptrs, mask=mask_m)
-        ds = (p * (dp - Di[:, None])) * sm_scale
-        ds = tl.where(p_mask, ds, 0.0).to(Q.dtype.element_ty)
-
-        # compute dk = dot(ds.T, q)
-        dk += tl.dot(tl.trans(ds), q)
-
-        # compute dq
-        if SEQUENCE_PARALLEL:
-            dq = tl.dot(ds, k)
-        else:
-            dq = tl.load(dq_ptrs, mask=q_mask, other=0.0)
-            dq += tl.dot(ds, k)
-        tl.store(dq_ptrs, dq.to(Q.dtype.element_ty), mask=q_mask)
-
-    # write-back dv and dk
-    dk_ptrs = dk_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk
-    dv_ptrs = dv_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vk
-
-    # write-back
-    tl.store(dk_ptrs, dk.to(K.dtype.element_ty), mask=kv_mask)
-    tl.store(dv_ptrs, dv.to(V.dtype.element_ty), mask=kv_mask)
-
-@triton.jit
-def _bwd_kernel(
-    Q,
-    K,
-    V,
-    sm_scale,
-    Out,
-    DO,
-    DQ,
-    DK,
-    DV,
-    L,
-    D,
-    stride_dq_all,
-    stride_qz,
-    stride_qh,
-    stride_qm,
-    stride_qk,
-    stride_kz,
-    stride_kh,
-    stride_kn,
-    stride_kk,
-    stride_vz,
-    stride_vh,
-    stride_vn,
-    stride_vk,
-    stride_deltaz,
-    stride_deltah,
-    stride_deltam,
-    Z,
-    H,
-    num_block_m,
-    num_block_n,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    max_seqlen_q,
-    max_seqlen_k,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    ACTUAL_BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    SEQUENCE_PARALLEL: tl.constexpr,
-    CAUSAL: tl.constexpr,
-    USE_EXP2: tl.constexpr,
-    IS_VARLEN: tl.constexpr,
-):
-    # program ids
-    off_hz = tl.program_id(0)
-    if SEQUENCE_PARALLEL:
-        start_n = tl.program_id(1)
-    off_z = off_hz // H
-    off_h = off_hz % H
-
-    if IS_VARLEN:
-        # Compute sequence lengths for the current batch
-        q_start = tl.load(cu_seqlens_q + off_z)
-        q_end = tl.load(cu_seqlens_q + off_z + 1)
-        k_start = tl.load(cu_seqlens_k + off_z)
-        k_end = tl.load(cu_seqlens_k + off_z + 1)
-
-        # Compute actual sequence lengths
-        N_CTX_Q = q_end - q_start
-        N_CTX_K = k_end - k_start
-    else:
-        q_start = 0
-        k_start = 0
-        N_CTX_Q = max_seqlen_q
-        N_CTX_K = max_seqlen_k
-
-    # input tensor offsets
-    q_offset = Q + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
-    k_offset = K + off_z * stride_kz + off_h * stride_kh + k_start * stride_kn
-    v_offset = V + off_z * stride_vz + off_h * stride_vh + k_start * stride_vn
-    do_offset = DO + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
-    l_offset = L + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam
-    d_offset = D + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam
-
-    # output tensor offsets
-    dk_offset = DK + off_z * stride_kz + off_h * stride_kh + k_start * stride_kn
-    dv_offset = DV + off_z * stride_vz + off_h * stride_vh + k_start * stride_vn
-    if SEQUENCE_PARALLEL:
-        dq_offset = DQ + start_n * stride_dq_all + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
-    else:
-        dq_offset = DQ + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
-
-    # inner loop
-    if SEQUENCE_PARALLEL:
-        _bwd_kernel_one_col_block(
-            Q,
-            K,
-            V,
-            sm_scale,
-            Out,
-            DO,
-            DQ,
-            DK,
-            DV,
-            L,
-            D,
-            q_offset,
-            k_offset,
-            v_offset,
-            do_offset,
-            dq_offset,
-            dk_offset,
-            dv_offset,
-            d_offset,
-            l_offset,
-            stride_dq_all,
-            stride_qz,
-            stride_qh,
-            stride_qm,
-            stride_qk,
-            stride_kz,
-            stride_kh,
-            stride_kn,
-            stride_kk,
-            stride_vz,
-            stride_vh,
-            stride_vn,
-            stride_vk,
-            stride_deltaz,
-            stride_deltah,
-            stride_deltam,
-            Z,
-            H,
-            N_CTX_Q,
-            N_CTX_K,
-            off_h,
-            off_z,
-            off_hz,
-            start_n,
-            num_block_m,
-            num_block_n,
-            BLOCK_M=BLOCK_M,
-            BLOCK_DMODEL=BLOCK_DMODEL,
-            ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
-            BLOCK_N=BLOCK_N,
-            SEQUENCE_PARALLEL=SEQUENCE_PARALLEL,
-            CAUSAL=CAUSAL,
-            USE_EXP2=USE_EXP2,
-        )
-    else:
-        for start_n in range(0, num_block_n):
-            _bwd_kernel_one_col_block(
-                Q,
-                K,
-                V,
-                sm_scale,
-                Out,
-                DO,
-                DQ,
-                DK,
-                DV,
-                L,
-                D,
-                q_offset,
-                k_offset,
-                v_offset,
-                do_offset,
-                dq_offset,
-                dk_offset,
-                dv_offset,
-                d_offset,
-                l_offset,
-                stride_dq_all,
-                stride_qz,
-                stride_qh,
-                stride_qm,
-                stride_qk,
-                stride_kz,
-                stride_kh,
-                stride_kn,
-                stride_kk,
-                stride_vz,
-                stride_vh,
-                stride_vn,
-                stride_vk,
-                stride_deltaz,
-                stride_deltah,
-                stride_deltam,
-                Z,
-                H,
-                N_CTX_Q,
-                N_CTX_K,
-                off_h,
-                off_z,
-                off_hz,
-                start_n,
-                num_block_m,
-                num_block_n,
-                BLOCK_M=BLOCK_M,
-                BLOCK_DMODEL=BLOCK_DMODEL,
-                ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
-                BLOCK_N=BLOCK_N,
-                SEQUENCE_PARALLEL=SEQUENCE_PARALLEL,
-                CAUSAL=CAUSAL,
-                USE_EXP2=USE_EXP2,
-            )
-
-
-# NOTE: smaller blocks have lower accuracy. more accumlation error probably 128 * 128 seems good but leads to oom. 64 * 64 has accumlation errors but no oom.
-def attention_prefill_backward_triton_impl(
-    do,
-    q,
-    k,
-    v,
-    o,
-    softmax_lse,
-    dq,
-    dk,
-    dv,
-    sm_scale: float,
-    alibi_slopes, # pylint: disable=unused-argument
-    causal,
-    layout: str,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    max_seqlen_q: int,
-    max_seqlen_k: int,
-    use_exp2: bool,
-    sequence_parallel = True,
-):
-    # make contigious
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    softmax_lse = softmax_lse.contiguous()
-
-    # get strides and shape
-    batch, nheads_q, nheads_k, head_size, max_seqlen_q, max_seqlen_k = get_shape_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) # pylint: disable=unused-variable
-    q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(q, k, v, o, layout)
-    stride_qz, stride_qh, stride_qm, stride_qk =  q_strides
-    stride_kz, stride_kh, stride_kn, stride_kk = k_strides
-    stride_vz, stride_vh, stride_vn, stride_vk = v_strides
-    stride_oz, stride_oh, stride_om, stride_ok = o_strides
-    batch_headsize = batch * nheads_q
-    is_varlen = layout == "thd"
-
-    # FIXME: some configs lead to oom for some reason when using 64 x 64 blocks
-    if max_seqlen_q <= 32 or max_seqlen_k <= 32:
-        BLOCK_M = 32
-        BLOCK_N = 32
-    else:
-        BLOCK_M = 64
-        BLOCK_N = 64
-    num_warps = 4 # NOTE: originial is 8. changing it to 1 caused issues be careful
-    num_stages = 1
-    waves_per_eu = 1
-
-    # divide up the problem
-    num_blocks_m = triton.cdiv(max_seqlen_q, BLOCK_M)
-    num_blocks_n = triton.cdiv(max_seqlen_k, BLOCK_N)
-
-    # get closest power of 2 over or equal to 32.
-    padded_d_model = 1 << (head_size - 1).bit_length()
-    padded_d_model = max(padded_d_model, 16)
-    BLOCK_DMODEL = padded_d_model
-    ACTUAL_BLOCK_DMODEL = head_size
-
-    do = do.contiguous()
-    # NOTE: we might need to copy the output tensor if they are not continuous or have other issues
-    copy_back = {"dq": False, "dk": False, "dv": False}
-
-    dq_og = None
-    # deal with dq
-    if dq is None:
-        if sequence_parallel:
-            dq = torch.zeros((num_blocks_n,) + q.shape, device=q.device, dtype=q.dtype)
-        else:
-            dq = torch.zeros(q.shape, device=q.device, dtype=q.dtype)
-    else:
-        dq_og = dq
-        if not dq.is_contiguous():
-            dq = dq.contiguous()
-            copy_back["dq"] = True
-
-        if sequence_parallel:
-            dq = torch.zeros((num_blocks_n,) + q.shape, device=q.device, dtype=q.dtype)
-            copy_back["dq"] = True
-        else:
-            # NOTE: the kernel does inplace accumlation so dq has to be zeros. This avoids the case where we are passed empty dq and it is not all zeros
-            dq.zero_()
-    stride_dq_all = dq.stride()[0]
-
-    dk_og = None
-    dv_og = None
-    # deal with dk, dv
-    if (dk is None) or (dv is None):
-        dk = torch.empty_like(k)
-        dv = torch.empty_like(v)
-    else:
-        if not dk.is_contiguous():
-            dk_og = dk
-            dk = dk.contiguous()
-            copy_back["dk"] = True
-
-        if not dv.is_contiguous():
-            dv_og = dv
-            dv = dv.contiguous()
-            copy_back["dv"] = True
-
-    # assert contigious
-    assert do.is_contiguous()
-    assert q.is_contiguous()
-    assert k.is_contiguous()
-    assert v.is_contiguous()
-    assert o.is_contiguous()
-    assert softmax_lse.is_contiguous()
-
-    # init delta
-    delta = torch.empty_like(softmax_lse)
-    if is_varlen:
-        stride_deltam, stride_deltah = delta.stride()
-        stride_deltaz = 0
-    else:
-        stride_deltaz, stride_deltah, stride_deltam = delta.stride()
-
-    _bwd_preprocess_use_o[(num_blocks_m, batch_headsize)](
-        o,
-        do,
-        delta,
-        stride_oz, stride_oh, stride_om, stride_ok,
-        stride_oz, stride_oh, stride_om, stride_ok,
-        stride_deltaz, stride_deltah, stride_deltam,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        max_seqlen_q,
-        max_seqlen_k,
-        BLOCK_M=BLOCK_M,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
-        N_CTX_Q=max_seqlen_q,
-        Z=batch,
-        H=nheads_q,
-        IS_VARLEN=is_varlen
-    )
-
-    _bwd_kernel[(batch_headsize, num_blocks_n if sequence_parallel else 1)](
-        q,
-        k,
-        v,
-        sm_scale,
-        o,
-        do,
-        dq,
-        dk,
-        dv,
-        softmax_lse,
-        delta,
-        stride_dq_all,
-        stride_qz, stride_qh, stride_qm, stride_qk,
-        stride_kz, stride_kh, stride_kn, stride_kk,
-        stride_vz, stride_vh, stride_vn, stride_vk,
-        stride_deltaz, stride_deltah, stride_deltam,
-        batch,
-        nheads_q,
-        num_blocks_m,
-        num_blocks_n,
-        cu_seqlens_q,
-        cu_seqlens_k,
-        max_seqlen_q,
-        max_seqlen_k,
-        BLOCK_M=BLOCK_M,
-        BLOCK_N=BLOCK_N,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
-        SEQUENCE_PARALLEL=sequence_parallel,
-        CAUSAL=causal,
-        USE_EXP2=use_exp2,
-        num_warps=num_warps,
-        num_stages=num_stages,
-        waves_per_eu = waves_per_eu,
-        IS_VARLEN=is_varlen
-    )
-
-    if sequence_parallel:
-        dq = dq.sum(dim=0)
-
-    if copy_back["dq"]:
-        dq_og.copy_(dq)
-        dq = dq_og
-    if copy_back["dk"]:
-        dk_og.copy_(dk)
-        dk = dk_og
-    if copy_back["dv"]:
-        dv_og.copy_(dv)
-        dv = dv_og
-
-    return dq, dk, dv, delta, None, None

modules/flash_attn_triton_amd/fwd_decode.py

@@ -1,700 +0,0 @@
-import torch
-import triton
-import triton.language as tl
-from modules.flash_attn_triton_amd.utils import _strides, get_padded_headsize
-
-
-@triton.jit
-def _fwd_kernel_splitK(
-    Q,
-    K,
-    V,
-    sm_scale,
-    Out_splitK,  # [B, H, split_k, Mq, K]
-    Metadata,  # [B, H, 2, split_k, M_ceil] contains [mi, li]
-    K_new,
-    V_new,
-    Cache_seqlens,
-    Cache_batch_idx,
-    Alibi_slopes,
-    stride_qz,
-    stride_qm,
-    stride_qg,
-    stride_qh,
-    stride_qd,
-    stride_kz,
-    stride_kn,
-    stride_kg,
-    stride_kh,
-    stride_kd,
-    stride_vz,
-    stride_vn,
-    stride_vg,
-    stride_vh,
-    stride_vd,
-    stride_osk_zhg,
-    stride_osk_s,
-    stride_osk_m,
-    stride_osk_d, # pylint: disable=unused-argument
-    stride_mzhg,
-    stride_m2,
-    stride_ms,
-    stride_mm, # pylint: disable=unused-argument
-    stride_kn_z,
-    stride_kn_n,
-    stride_kn_g,
-    stride_kn_h,
-    stride_kn_d,
-    stride_vn_z,
-    stride_vn_n,
-    stride_vn_g,
-    stride_vn_h,
-    stride_vn_d,
-    stride_az,
-    stride_ah,
-    Z, # pylint: disable=unused-argument
-    N_CTX_Q,
-    N_CTX_K,
-    N_CTX_NEW,
-    BLOCK_N_PER_SPLIT,
-    H_q: tl.constexpr,
-    H_kv: tl.constexpr,
-    G_q: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_DMODEL: tl.constexpr,
-    ACTUAL_BLOCK_DMODEL: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    BOUNDS_CHECKS_N: tl.constexpr,
-    USE_CACHE_SEQLENs: tl.constexpr,
-    USE_CACHE_BATCH_IDX: tl.constexpr,
-    NEW_KV: tl.constexpr,
-    IS_GQA: tl.constexpr,
-    IS_CAUSAL: tl.constexpr,
-    USE_ALIBI: tl.constexpr,
-):
-    # Padding
-    PADDED_HEAD: tl.constexpr = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
-    if PADDED_HEAD:
-        d_mask = tl.arange(0, BLOCK_DMODEL) < ACTUAL_BLOCK_DMODEL
-
-    start_m = tl.program_id(0)
-    off_zhg = tl.program_id(1)
-    off_z = off_zhg // (H_q * G_q)
-    off_h_q = (off_zhg // G_q) % H_q
-    off_g_q = off_zhg % G_q
-    splitk_idx = tl.program_id(2)
-
-    # pick batch index
-    if USE_CACHE_BATCH_IDX:
-        cache_batch_idx = tl.load(Cache_batch_idx + off_z)
-    else:
-        cache_batch_idx = off_z
-
-    # Load ALiBi slope if enabled
-    if USE_ALIBI:
-        a_offset = off_z * stride_az + off_h_q * stride_ah
-        alibi_slope = tl.load(Alibi_slopes + a_offset)
-    else:
-        alibi_slope = None
-
-    lo = splitk_idx * BLOCK_N_PER_SPLIT
-    if USE_CACHE_SEQLENs:
-        cache_seqlen_last_idx = tl.load(Cache_seqlens + off_z)
-        if NEW_KV:
-            kv_len = cache_seqlen_last_idx + N_CTX_NEW
-        else:
-            kv_len = cache_seqlen_last_idx
-    else:
-        kv_len = N_CTX_K
-    hi = tl.minimum((splitk_idx + 1) * BLOCK_N_PER_SPLIT, kv_len)
-
-    HEAD_RATIO: tl.constexpr = H_q // H_kv
-    if IS_GQA:
-        k_head_idx = off_h_q // HEAD_RATIO
-        v_head_idx = k_head_idx
-    else:
-        k_head_idx = off_h_q
-        v_head_idx = off_h_q
-
-    # calculate base offset
-    k_base = K + k_head_idx * stride_kh + cache_batch_idx * stride_kz + off_g_q * stride_kg
-    v_base = V + v_head_idx * stride_vh + cache_batch_idx * stride_vz + off_g_q * stride_vg
-
-    # Copy new Keys and Values into Cache
-    if NEW_KV:
-        knew_base = K_new + k_head_idx * stride_kn_h + off_z * stride_kn_z + off_g_q * stride_kn_g
-
-        # Determine the starting position for new data in the cache
-        if USE_CACHE_SEQLENs:
-            start_idx = tl.load(Cache_seqlens + off_z)
-        else:
-            start_idx = N_CTX_K - N_CTX_NEW
-
-        # Copy new Keys
-        for i in range(0, N_CTX_NEW, BLOCK_N):
-            # Load from K_new
-            k_new_block = tl.load(
-                knew_base +
-                tl.arange(0, BLOCK_DMODEL)[:, None] * stride_kn_d +
-                (tl.arange(0, BLOCK_N) + i)[None, :] * stride_kn_n,
-                 mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) &
-                     (tl.arange(0, BLOCK_DMODEL)[:, None] < ACTUAL_BLOCK_DMODEL),
-                other=0
-            )
-
-            # Store to K
-            tl.store(
-                k_base +
-                tl.arange(0, BLOCK_DMODEL)[:, None] * stride_kd +
-                (tl.arange(0, BLOCK_N) + i + start_idx)[None, :] * stride_kn,
-                k_new_block,
-                 mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) &
-                     (tl.arange(0, BLOCK_DMODEL)[:, None] < ACTUAL_BLOCK_DMODEL),
-            )
-
-        # Copy new Values
-        vnew_base = V_new + v_head_idx * stride_vn_h + off_z * stride_vn_z + off_g_q * stride_vn_g
-        for i in range(0, N_CTX_NEW, BLOCK_N):
-            # Load from V_new
-            v_new_block = tl.load(
-                vnew_base +
-                (tl.arange(0, BLOCK_N) + i)[:, None] * stride_vn_n +
-                tl.arange(0, BLOCK_DMODEL)[None, :] * stride_vn_d,
-                mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) &
-                     (tl.arange(0, BLOCK_DMODEL)[None, :] < ACTUAL_BLOCK_DMODEL),
-                other=0
-            )
-
-            # Store to V
-            tl.store(
-                v_base +
-                (tl.arange(0, BLOCK_N) + i + start_idx)[:, None] * stride_vn +
-                tl.arange(0, BLOCK_DMODEL)[None, :] * stride_vd,
-                v_new_block,
-                 mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) &
-                     (tl.arange(0, BLOCK_DMODEL)[None, :] < ACTUAL_BLOCK_DMODEL),
-            )
-
-    Q_block_ptr = tl.make_block_ptr(
-        base=Q + off_h_q * stride_qh + off_z * stride_qz + off_g_q * stride_qg,
-        shape=(N_CTX_Q, ACTUAL_BLOCK_DMODEL),
-        strides=(stride_qm, stride_qd),
-        offsets=(start_m * BLOCK_M, 0),
-        block_shape=(BLOCK_M, BLOCK_DMODEL),
-        order=(1, 0),
-    )
-
-    K_block_ptr = tl.make_block_ptr(
-        base=k_base,
-        shape=(ACTUAL_BLOCK_DMODEL, hi),
-        strides=(stride_kd, stride_kn),
-        offsets=(0, lo),
-        block_shape=(BLOCK_DMODEL, BLOCK_N),
-        order=(0, 1),
-    )
-    V_block_ptr = tl.make_block_ptr(
-        base=v_base,
-        shape=(hi, ACTUAL_BLOCK_DMODEL),
-        strides=(stride_vn, stride_vd),
-        offsets=(lo, 0),
-        block_shape=(BLOCK_N, BLOCK_DMODEL),
-        order=(1, 0),
-    )
-
-    K_scale_shift_block_ptr = None
-    V_scale_shift_block_ptr = None
-
-    # initialize pointer to m and l
-    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
-    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-
-    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)  # noqa: F821
-
-    # scale sm_scale by log_2(e) and use
-    # 2^x instead of exp in the loop because CSE and LICM
-    # don't work as expected with `exp` in the loop
-    qk_scale = sm_scale * 1.44269504
-    # load q: it will stay in SRAM throughout
-    q = tl.load(  # noqa: F821
-        tl.advance(Q_block_ptr, (0, 0)), boundary_check=(0, ))
-    q = (q * qk_scale).to(q.dtype)
-    if PADDED_HEAD:
-        q = tl.where(d_mask[None, :], q, 0.0)
-
-    # loop over k, v and update accumulator
-    for start_n in range(lo, hi, BLOCK_N):
-        k, v = load_k_v_group(
-            K_block_ptr,
-            V_block_ptr,
-            K_scale_shift_block_ptr,
-            V_scale_shift_block_ptr,
-            BOUNDS_CHECKS_N,
-            1,
-            BLOCK_DMODEL,
-            ACTUAL_BLOCK_DMODEL,
-            Q.dtype.element_ty,
-            0,
-        )
-        if PADDED_HEAD:
-            k = tl.where(d_mask[:, None], k, 0.0)
-            v = tl.where(d_mask[None, :], v, 0.0)
-
-        # -- compute qk ---
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)  # noqa: F821
-
-        if USE_ALIBI:
-            row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-            col_idx = start_n + tl.arange(0, BLOCK_N)
-
-            # Compute relative positions
-            relative_pos = row_idx[:, None] + kv_len - (N_CTX_Q + col_idx[None, :])
-            relative_pos = tl.abs(relative_pos)
-
-            # Compute ALiBi bias
-            alibi_bias = -1 * alibi_slope * relative_pos
-            qk += (alibi_bias * 1.44269504)
-
-        # Apply causal mask if IS_CAUSAL is True
-        if IS_CAUSAL:
-            row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-            col_idx = start_n + tl.arange(0, BLOCK_N)
-
-            # create a N_CTX_Q x kv_len causal mask
-            col_offset = N_CTX_Q - kv_len
-            causal_mask = row_idx[:, None] >= (col_offset + col_idx[None, :])
-
-            # Apply the mask
-            qk = tl.where(causal_mask, qk, float("-inf"))
-
-        # TODO: This is slow, and only needed at the last iteration.
-        # Maybe we can unroll the last iteration instead?
-        if BOUNDS_CHECKS_N:
-            qk = tl.where(tl.arange(0, BLOCK_N) < hi - start_n, qk, float("-inf"))
-
-        # -- compute scaling constant ---
-        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
-        if IS_CAUSAL:
-            alpha = tl.math.exp2(tl.where(m_i > float("-inf"), m_i - m_i_new, float("-inf")))
-        else:
-            alpha = tl.math.exp2(m_i - m_i_new)
-        # cause of nan because subtracting infs
-        if IS_CAUSAL:
-            qk = tl.where(qk > float("-inf"), qk - m_i_new[:, None], float("-inf"))
-        else:
-            qk = qk - m_i_new[:, None]
-
-        p = tl.math.exp2(qk)
-
-        # -- update m_i and l_i --
-        l_i = l_i * alpha + tl.sum(p, 1)
-        m_i = m_i_new
-        p = p.to(Q.dtype.element_ty)
-
-        # -- scale and update acc --
-        acc *= alpha[:, None]
-        acc += tl.dot(p.to(v.dtype), v)
-
-        # update pointers
-        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
-        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
-
-    # write back O
-    O_block_ptr = tl.make_block_ptr(
-        base=Out_splitK + off_zhg * stride_osk_zhg + splitk_idx * stride_osk_s,
-        shape=(N_CTX_Q, BLOCK_DMODEL),
-        strides=(stride_osk_m, 1),
-        offsets=(start_m * BLOCK_M, 0),
-        block_shape=(BLOCK_M, BLOCK_DMODEL),
-        order=(1, 0),
-    )
-    tl.store(
-        tl.advance(O_block_ptr, (0, 0)),
-        acc,
-        boundary_check=(0, ),
-    )
-    # Write metadata for split-K reduction
-    Metadata_ptr = (Metadata + off_zhg * stride_mzhg + splitk_idx * stride_ms + start_m * BLOCK_M +
-                    tl.arange(0, BLOCK_M))
-    tl.store(Metadata_ptr, m_i)
-    tl.store(Metadata_ptr + stride_m2, l_i)
-
-
-@triton.jit
-def load_k_v_group(
-    K_block_ptr,
-    V_block_ptr,
-    K_scale_shift_block_ptr, V_scale_shift_block_ptr, # pylint: disable=unused-argument
-    BOUNDS_CHECKS_N: tl.constexpr,
-    PACKED_PER_VAL: tl.constexpr, BLOCK_DMODEL: tl.constexpr, # pylint: disable=unused-argument
-    ACTUAL_BLOCK_DMODEL: tl.constexpr,
-    dtype: tl.constexpr, # pylint: disable=unused-argument
-    group_id: tl.constexpr,
-):
-    # Load K/V for a given block
-    # Advance to the current quantization group
-    K_block_ptr = tl.advance(K_block_ptr, (ACTUAL_BLOCK_DMODEL * group_id, 0))
-    V_block_ptr = tl.advance(V_block_ptr, (0, ACTUAL_BLOCK_DMODEL * group_id))
-
-    # -- load k, v --
-    k = tl.load(K_block_ptr, boundary_check=(1, ) if BOUNDS_CHECKS_N else ())
-    v = tl.load(V_block_ptr, boundary_check=(0, ) if BOUNDS_CHECKS_N else ())
-
-    return k, v
-
-
-@triton.jit
-def cast_uint32_to_half2(scale_shift):
-    # Extract two float16 packed into one int32
-    scale = scale_shift & 0xFFFF
-    shift = scale_shift >> 16
-    scale = scale.to(tl.uint16).to(tl.float16, bitcast=True)
-    shift = shift.to(tl.uint16).to(tl.float16, bitcast=True)
-    return scale, shift
-
-
-@triton.jit
-def dequantize(
-    x_,
-    scale,
-    shift,
-    PACKED_PER_VAL: tl.constexpr = 8,
-):
-    # PACKED_PER_VAL is the number of values packed into
-    # each element x_. For example, for int4 quantization
-    #and x_ of type int32, PACKED_PER_VAL is 8.
-
-    BLOCK_N: tl.constexpr = x_.shape[0]
-    BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1]
-    offsets = tl.arange(0, PACKED_PER_VAL) * 4
-    quant_offset = (x_[:, None, :] >> offsets[None, :, None])  # (BLOCK_N, PACKED_PER_VAL, D // PACKED_PER_VAL)
-
-    quant_offset = tl.view(quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL))
-    # Trick - instead of converting int4 to float16 we view it as float16
-    # and then multiply by 32768 * 512 == 2**24
-    quant_offset = (quant_offset & 0xF).to(tl.uint16).to(tl.float16, bitcast=True)
-    quant_offset = (quant_offset * 32768.0).to(tl.float16)
-    scale_512 = scale * 512
-
-    dequant = quant_offset * scale_512 + shift
-    return dequant
-
-
-@triton.jit
-def _splitK_reduce(
-    Out_splitK,  # [B, H, split_k, Mq, K]
-    Metadata,  # [B, H, 2, split_k, M_ceil] contains [mi, li]
-    Out,  # [B, H, M, K]
-    LSE,  # [B, H, M]
-    stride_osk_zhg,
-    stride_osk_s,
-    stride_osk_m,
-    stride_osk_k,
-    stride_mzhg,
-    stride_m2,
-    stride_ms,
-    stride_mm,
-    stride_oz,
-    stride_oh,
-    stride_og,
-    stride_om,
-    stride_ok, # pylint: disable=unused-argument
-    stride_lse_zhg,
-    stride_lse_m, M_ceil: tl.constexpr, # pylint: disable=unused-argument
-    BLOCK_SIZE: tl.constexpr,
-    H: tl.constexpr,
-    G: tl.constexpr,
-    split_k: tl.constexpr,
-    splitK_pow2: tl.constexpr,
-    use_mask: tl.constexpr,
-    IS_CAUSAL: tl.constexpr,
-):
-    off_zhg = tl.program_id(0)
-    off_z = off_zhg // (H * G)
-    off_h = (off_zhg // G) % H
-    off_g = off_zhg % G
-    off_m = tl.program_id(1)
-    off_k = tl.program_id(2)
-
-    # read  chunk
-    spk_idx = tl.arange(0, splitK_pow2)
-    kidx = tl.arange(0, BLOCK_SIZE)
-
-    Metadata_ptr = Metadata + stride_mzhg * off_zhg + spk_idx * stride_ms + off_m * stride_mm
-
-    o_ptr = (Out_splitK + off_zhg * stride_osk_zhg + stride_osk_m * off_m + off_k * BLOCK_SIZE +
-             stride_osk_s * spk_idx[:, None] + kidx[None, :] * stride_osk_k)
-
-    # read max values of each splitK
-    if use_mask:
-        spk_mask = spk_idx < split_k
-        l_m = tl.load(Metadata_ptr, mask=spk_mask, other=float("-inf"))
-        l_sum = tl.load(Metadata_ptr + stride_m2, mask=spk_mask, other=0.0)
-        acc = tl.load(o_ptr, mask=spk_mask[:, None], other=0.0)
-    else:
-        l_m = tl.load(Metadata_ptr)
-        l_sum = tl.load(Metadata_ptr + stride_m2)
-        acc = tl.load(o_ptr)
-
-    g_m = tl.max(l_m, axis=0)
-
-    if IS_CAUSAL:
-        l_m_offset = l_m - g_m
-        alpha = tl.where(l_m_offset > float("-inf"), tl.math.exp2(l_m_offset), 0.0)
-    else:
-        alpha = tl.math.exp2(l_m - g_m)
-
-    # read sum
-    l_sum *= alpha
-    g_sum = tl.sum(l_sum, axis=0)
-    acc = acc * alpha[:, None]
-
-    if IS_CAUSAL:
-        # Avoid division by zero
-        g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0)
-        acc_out = tl.sum(acc, axis=0) / g_sum_safe
-    else:
-        acc_out = tl.sum(acc, axis=0) / g_sum
-
-    # Store output
-    Out_ptr = (Out + stride_oz * off_z + stride_oh * off_h + stride_og * off_g + stride_om * off_m +
-               off_k * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE))
-    tl.store(Out_ptr, acc_out)
-
-    # Store lse
-    l_ptrs = LSE + off_zhg * stride_lse_zhg + off_m
-    if IS_CAUSAL:
-        lse = tl.where(g_sum > 0, (g_m + tl.math.log2(g_sum)) / 1.44269504, g_m)
-        tl.store(l_ptrs, lse)
-    else:
-        tl.store(l_ptrs, (g_m + tl.math.log2(g_sum)) / 1.44269504)
-
-
-def quantize_kv_int4(k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
-    # Scale and shift are such that quantization linearly maps
-    # int4 values range [0..15] to input values range min(k)..max(k)
-    # individually for every row
-    k = k.reshape(*k.shape[:-1], num_groups, k.shape[-1] // num_groups)
-    max_vals = torch.max(k, dim=-1, keepdim=True).values
-    min_vals = torch.min(k, dim=-1, keepdim=True).values
-    scale_k: torch.Tensor = (max_vals - min_vals) / 15
-
-    shift_k = torch.min(k, dim=-1, keepdim=True).values
-    scale_k = scale_k.to(torch.float16)
-    shift_k = shift_k.to(torch.float16)
-
-    in_bytes = ((k - shift_k.expand(k.shape)) / scale_k.expand(k.shape)) + 0.5
-    in_bytes = in_bytes.to(torch.uint8)
-    in_int4 = in_bytes & 0xF
-    in_int4_packed = in_int4[..., ::2] + (in_int4[..., 1::2] << 4)
-    scale_shift = torch.concat([scale_k.view(torch.uint8), shift_k.view(torch.uint8)], dim=-1)
-    k_quant = torch.concat(
-        [
-            scale_shift.flatten(start_dim=-2),
-            in_int4_packed.flatten(start_dim=-2),
-        ],
-        dim=-1,
-    ).view(torch.int16)
-    return k_quant
-
-
-def dequantize_kv_fp16(quant_k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
-    k_i16 = quant_k.view(torch.int16)
-    k_ui8 = k_i16.view(torch.uint8)
-
-    ss_size = num_groups * 4
-    scale_shift_ui8 = k_ui8[..., 0:ss_size]
-    scale_shift_ui8 = scale_shift_ui8.reshape(*scale_shift_ui8.shape[:-1], num_groups, 4)
-    scale = scale_shift_ui8[..., 0:2].view(torch.float16)
-    shift = scale_shift_ui8[..., 2:4].view(torch.float16)
-
-    kv_ui8 = k_ui8[..., ss_size:]
-    k_ui8 = kv_ui8.reshape(*kv_ui8.shape[:-1], num_groups, -1)
-    k1_i4 = k_ui8 & 0xF
-    k2_i4 = (k_ui8 & 0xF0) >> 4
-    k_shape = k1_i4.shape
-    k1_f16 = k1_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape)
-    k2_f16 = k2_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape)
-
-    out = torch.empty((*k1_f16.shape[:-1], k1_f16.shape[-1] * 2), dtype=torch.float16, device=quant_k.device)
-    out[..., ::2] = k1_f16
-    out[..., 1::2] = k2_f16
-    out = out.reshape(*k_shape[:-2], -1)
-
-    return out
-
-
-def get_split_k(B: int, G: int, H: int, Mk: int) -> int:
-    """Heuristic for the number of splits"""
-    bh = max(B * H, 1)  # NOTE: Handle B*h=0 case
-    split_k = max(Mk, 1024) // bh
-    max_chunk_size = 64
-    while split_k > 0 and Mk / split_k < max_chunk_size:
-        split_k = split_k // 2
-    while B * H * G * split_k >= 1024:
-        split_k = split_k // 2
-    split_k = min(split_k, 512)
-    split_k = max(split_k, 1)
-    return split_k
-
-def attention_decode_forward_triton_impl(q, k, v, sm_scale, causal, alibi_slopes, layout, cache_seqlens, cache_batch_idx, new_kv, k_new, v_new):
-    # kernel config
-    BLOCK_M = 16
-    BLOCK_N = 64
-    SPLIT_K = None
-    NUM_QUANT_GROUPS = 1 # pylint: disable=unused-variable
-
-    # kernels expects "bsghd"
-    original_layout = layout
-    if layout == "bshd":
-        q = q.unsqueeze(2)
-        k = k.unsqueeze(2)
-        v = v.unsqueeze(2)
-        if new_kv:
-            k_new = k_new.unsqueeze(2)
-            v_new = v_new.unsqueeze(2)
-        layout = "bsghd"
-    elif layout == "bhsd":
-        q = q.permute(0, 2, 1, 3).unsqueeze(2)
-        k = k.permute(0, 2, 1, 3).unsqueeze(2)
-        v = v.permute(0, 2, 1, 3).unsqueeze(2)
-        if new_kv:
-            k_new = k_new.permute(0, 2, 1, 3).unsqueeze(2)
-            v_new = v_new.permute(0, 2, 1, 3).unsqueeze(2)
-        layout = "bsghd"
-    elif layout == "bsghd":
-        pass
-    elif layout is None:
-        raise ValueError("Layout not given")
-    assert layout == "bsghd"
-
-    # get dims
-    batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_q = q.shape
-    _, seqlen_k, n_group_k, heads_per_group_k, dim_k = k.shape # pylint: disable=unused-variable
-    _, seqlen_v, n_group_v, heads_per_group_v, dim_v = v.shape # pylint: disable=unused-variable
-
-    assert dim_q == dim_k == dim_v, f"Dimensions must match: {dim_q}, {dim_k}, {dim_v}"
-
-    # get padded size
-    dim_padded  = get_padded_headsize(dim_k)
-
-    # Handle MQA/GQA case
-    if heads_per_group_q > heads_per_group_k:
-        is_gqa = True
-    elif heads_per_group_q < heads_per_group_k:
-        raise ValueError("heads_per_group_q < heads_per_group_k")
-    else:
-        is_gqa = False
-
-    assert dim_k == dim_q, f"Keys have head dim {dim_k} but queries have head dim {dim_q}"
-
-    if SPLIT_K is not None:
-        split_k = SPLIT_K
-    else:
-        # Use heuristics
-        split_k = get_split_k(batch_size, n_group_q, heads_per_group_q, seqlen_k) # NOTE: should the split think about seqlens?
-
-    seqlen_q_ceil = (seqlen_q + BLOCK_M - 1) // BLOCK_M * BLOCK_M
-    out_splitk = torch.empty([batch_size * n_group_q * heads_per_group_q, split_k, seqlen_q_ceil, dim_padded], dtype=torch.float32, device=q.device)
-    metadata = torch.empty([batch_size * n_group_q * heads_per_group_q, 2, split_k, seqlen_q_ceil], dtype=torch.float32, device=q.device)
-    lse = torch.empty((batch_size * n_group_q * heads_per_group_q, seqlen_q), device=q.device, dtype=torch.float32)
-    grid = (triton.cdiv(seqlen_q, BLOCK_M), batch_size * n_group_q * heads_per_group_q, split_k)
-
-    num_warps = 1
-    split_size = (seqlen_k + split_k - 1) // split_k
-    use_cache_seqlens = cache_seqlens is not None
-
-    # TODO: enable quantization
-    _fwd_kernel_splitK[grid](
-        Q=q,
-        K=k,
-        V=v,
-        sm_scale=sm_scale,
-        Out_splitK=out_splitk,
-        Metadata=metadata,
-        K_new = k_new,
-        V_new = v_new,
-        Cache_seqlens=cache_seqlens,
-        Cache_batch_idx=cache_batch_idx,
-        Alibi_slopes=alibi_slopes,
-        **_strides(q, "qz", "qm", "qg", "qh", "qd"),
-        **_strides(k, "kz", "kn", "kg", "kh", "kd"),
-        **_strides(v, "vz", "vn", "vg", "vh", "vd"),
-        **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_d"),
-        **_strides(metadata, "mzhg", "m2", "ms", "mm"),
-        **_strides(k_new, "kn_z", "kn_n", "kn_g", "kn_h", "kn_d"),
-        **_strides(v_new, "vn_z", "vn_n", "vn_g", "vn_h", "vn_d"),
-        **_strides(alibi_slopes, "az", "ah"),
-        Z=batch_size,
-        H_q=heads_per_group_q,
-        H_kv=heads_per_group_k,
-        G_q=n_group_q,
-        N_CTX_Q=seqlen_q,
-        N_CTX_K=seqlen_k,
-        N_CTX_NEW=k_new.shape[1] if new_kv else None,
-        BLOCK_N_PER_SPLIT=split_size,
-        BLOCK_M=BLOCK_M,
-        BLOCK_N=BLOCK_N,
-        BLOCK_DMODEL=dim_padded,
-        ACTUAL_BLOCK_DMODEL=dim_k,
-        BOUNDS_CHECKS_N=(split_size % BLOCK_N) > 0 or use_cache_seqlens,
-        USE_CACHE_SEQLENs=use_cache_seqlens,
-        USE_CACHE_BATCH_IDX=cache_batch_idx is not None,
-        NEW_KV=new_kv,
-        IS_GQA=is_gqa,
-        IS_CAUSAL=causal,
-        USE_ALIBI=False if alibi_slopes is None else True,
-        num_warps=num_warps,
-        num_stages=1,
-    )
-
-    out = torch.empty((batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_padded), device=q.device, dtype=q.dtype)
-
-    # Merge together
-    splitK_pow2 = triton.next_power_of_2(split_k)
-    use_mask = splitK_pow2 > split_k
-    if batch_size * n_group_q * heads_per_group_q * seqlen_q >= 512:
-        k_block_num = 1
-    else:
-        k_block_num = 2
-    assert dim_padded % k_block_num == 0
-    k_block_size = dim_padded // k_block_num
-    grid = (batch_size * n_group_q * heads_per_group_q, seqlen_q, k_block_num)
-
-    _splitK_reduce[grid](
-        out_splitk,
-        metadata,
-        out,
-        lse,
-        **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_k"),
-        **_strides(metadata, "mzhg", "m2", "ms", "mm"),
-        **_strides(out, "oz", "om", "og", "oh", "ok"),
-        **_strides(lse, "lse_zhg", "lse_m"),
-        M_ceil=seqlen_q_ceil,
-        BLOCK_SIZE=k_block_size,
-        G=n_group_q,
-        H=heads_per_group_q,
-        # TODO: Tune num_warps
-        split_k=split_k,
-        splitK_pow2=splitK_pow2,
-        use_mask=use_mask,
-        IS_CAUSAL=causal,
-        num_warps=4)
-
-    lse = lse.reshape([batch_size, n_group_q, heads_per_group_q, seqlen_q])
-    if q.ndim == 4:
-        # BMGHK -> BMHK
-        assert n_group_q == 1
-        out = out[:, :, 0]
-        lse = lse[:, 0]
-    if seqlen_k == 0:
-        out.zero_()
-    out = out.reshape(batch_size, heads_per_group_q * n_group_q, -1, dim_padded).contiguous()
-
-    # output is batch_size, heads_per_group_q * group_q, seqlen_q, dim_q
-    if original_layout == "bshd":
-        # out=out.transpose(1, 2).contiguous() # this screws up heads and data.
-        # the data is laid out properly. Just need to reshape dims
-        out = out.reshape(batch_size, seqlen_q, -1, dim_padded)
-
-    return out.narrow(-1, 0, dim_k), lse

modules/flash_attn_triton_amd/fwd_prefill.py

@@ -1,33 +1,8 @@
+from typing import Literal, Optional, Union
 import torch
 import triton
 import triton.language as tl
-from modules.flash_attn_triton_amd.utils import get_shape_from_layout, get_strides_from_layout, is_cdna, is_rdna, AUTOTUNE
-
-
-@triton.jit
-def cdiv_fn(x, y):
-    return (x + y - 1) // y
-
-
-@triton.jit
-def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride): # pylint: disable=unused-argument
-    ms = tl.arange(0, m)
-    ns = tl.arange(0, n)
-    return philox_offset + ms[:, None] * stride + ns[None, :]
-
-
-@triton.jit
-def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):
-    rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride).to(tl.uint32)
-    # TODO: use tl.randint for better performance
-    return tl.rand(philox_seed, rng_offsets)
-
-
-@triton.jit
-def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):
-    rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride)
-    rng_keep = rng_output > dropout_p
-    return rng_keep
+from modules.flash_attn_triton_amd.utils import AUTOTUNE, compute_alibi_block, get_shapes_from_layout, get_strides_from_layout, is_cdna, is_rdna
 
 
 # Convenience function to load with optional boundary checks.
@@ -50,47 +25,14 @@ def load_fn(ptrs, offset_first, offset_second, boundary_first, boundary_second):
 
 
 @triton.jit
-def compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False):
-    # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix
-    # for casual mask we want something like this where (1 is kept and 0 is masked)
-    # seqlen_q = 2 and seqlen_k = 5
-    #   1 1 1 1 0
-    #   1 1 1 1 1
-    # seqlen_q = 5 and seqlen_k = 2
-    #        0 0
-    #        0 0
-    #        0 0
-    #        1 0
-    #        1 1
-    # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal
-    # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5, offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False
-    # 1. offs_m[:,None] = [[0],
-    #                       [1],
-    # 2. offs_m[:,None] + seqlen_k = [[5],
-    #                                  [6],
-    # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3],
-    #                                             [4],
-    # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] = [[3], - [[0, 1, 2, 3, 4]] =  [[ 3, 2, 1, 0,-1],
-    #                                                            [4],                           [ 4, 3, 2, 1, 0]]
-    # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1],
-    #                                                     [ -4, -3, -2, -1, 0]],
-    relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :]
-    alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
-    if transpose:
-        return alibi_block.T
-    else:
-        return alibi_block
-
-
-@triton.jit
-def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, start_m,
-                    actual_seqlen_k, actual_seqlen_q, dropout_p, philox_seed, batch_philox_offset, exp_scores_ptrs,
-                    block_min, block_max, offs_n_causal, masked_blocks, n_extra_tokens, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs, # pylint: disable=unused-argument
+def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, start_m,
+                    actual_seqlen_k, actual_seqlen_q, dropout_p, philox_seed, philox_ptrs, sd_mask_ptrs, dropout_mask_ptrs,
+                    block_min, block_max, offs_n_causal, masked_blocks, n_extra_tokens, alibi_slope, # pylint: disable=unused-argument
                     IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
                     OFFS_M: tl.constexpr, OFFS_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, MASK_STEPS: tl.constexpr,
                     ENABLE_DROPOUT: tl.constexpr, PADDED_HEAD: tl.constexpr,
-                    ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_EXP2: tl.constexpr,
-                    RETURN_SCORES: tl.constexpr):
+                    ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr,
+                    RETURN_SCORES: tl.constexpr, ACCUMULATOR_TYPE):
     if USE_EXP2:
         RCP_LN2: tl.constexpr = 1.4426950408889634
 
@@ -107,7 +49,7 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
         if PRE_LOAD_V:
             # We can use the same offsets as k, just with dims transposed.
             v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL)
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=ACCUMULATOR_TYPE)
         # We start from end of seqlen_k so only the first iteration would need
         # to be checked for padding if it is not a multiple of block_n
         # TODO: This can be optimized to only be true for the padded block.
@@ -117,18 +59,20 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
             # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not is_modulo_mn.
             # last step might get wasted but that is okay. check if this masking works For
             # that case.
-            if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):
+            if start_n + BLOCK_N == block_max and n_extra_tokens != 0:
                 boundary_m = tl.full([BLOCK_M], actual_seqlen_k, dtype=tl.int32)
                 size_n = start_n + OFFS_N[None, :]
                 mask = size_n < boundary_m[:, None]
                 qk = tl.where(mask, qk, float("-inf"))
 
+        # compute masks
+        q_mask = OFFS_M[:, None] < actual_seqlen_q
+        k_mask = (start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k
+        p_mask = q_mask & k_mask
+
         # -- compute qk ----
         qk += tl.dot(q, k)
-        qk_scaled =  qk * SM_SCALE
-        if RETURN_SCORES:
-            score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
-            tl.store(score_ptrs, qk_scaled, mask=score_mask)
+        qk_scaled = qk * SM_SCALE
 
         if IS_CAUSAL:
             causal_boundary = start_n + offs_n_causal
@@ -139,8 +83,8 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
             bias = load_fn(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q, actual_seqlen_k)
             qk_scaled += bias
 
-        if alibi_slope is not None:
-            # Compute the global position of each token within the sequence
+        if USE_ALIBI:
+            # compute the global position of each token within the sequence
             global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
             global_n_positions = start_n + tl.arange(0, BLOCK_N)
             alibi_block = compute_alibi_block(alibi_slope, actual_seqlen_q, actual_seqlen_k, global_m_positions,
@@ -151,10 +95,6 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
 
         # scale and subtract max
         q_shifted = qk_scaled - m_ij[:, None]
-        if RETURN_SCORES:
-            # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
-            scores_scaled_shifted_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
-            tl.store(scores_scaled_shifted_ptrs, q_shifted, mask=scores_scaled_shifted_mask)
 
         # Compute scaled QK and softmax probabilities
         if USE_EXP2:
@@ -165,17 +105,18 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
         # CAVEAT: Must update l_ij before applying dropout
         l_ij = tl.sum(p, 1)
         if ENABLE_DROPOUT:
-            philox_offset = batch_philox_offset + start_m * BLOCK_M * actual_seqlen_k + start_n - BLOCK_N
-            keep = dropout_mask(philox_seed, philox_offset, dropout_p, BLOCK_M, BLOCK_N, actual_seqlen_k)
-            if RETURN_SCORES:
-                # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
-                exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
-                tl.store(exp_scores_ptrs, tl.where(keep, p, -p), mask=exp_score_mask)
-            p = tl.where(keep, p, 0.0)
+            rng_output = tl.rand(philox_seed, philox_ptrs)  # TODO: use tl.randint for better performance
+            dropout_mask = rng_output > dropout_p
+
+            # return scores with negative values for dropped vals
+            sd_mask = tl.where(dropout_mask, p, -p)
+            tl.store(sd_mask_ptrs, sd_mask, mask=p_mask)
+
+            # apply dropout mask in place
+            p = tl.where(dropout_mask, p, 0.0)
         elif RETURN_SCORES:
             # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
-            exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
-            tl.store(exp_scores_ptrs, p, mask=exp_score_mask)
+            tl.store(sd_mask_ptrs, p, mask=p_mask)
 
         # -- update output accumulator --
         # alpha is an adjustment factor for acc and li as we loop and find new maxes
@@ -186,7 +127,6 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
         else:
             alpha = tl.math.exp(m_diff)
         acc = acc * alpha[:, None]
-        v = None
         if not PRE_LOAD_V:
             v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL)
         # -- update m_i and l_i
@@ -199,9 +139,11 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri
         if bias_ptrs is not None:
             bias_ptrs += BLOCK_N * stride_bn
         if RETURN_SCORES:
-            score_ptrs += BLOCK_N
-            scores_scaled_shifted_ptrs += BLOCK_N
-            exp_scores_ptrs += BLOCK_N
+            sd_mask_ptrs += BLOCK_N * stride_sn
+
+        if ENABLE_DROPOUT:
+            dropout_mask_ptrs += BLOCK_N * stride_sn
+            philox_ptrs += BLOCK_N * stride_sn
     return acc, l_i, m_i
 
 
@@ -222,7 +164,7 @@ def get_cdna_autotune_configs():
         # Fall-back config.
         triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1,
                       num_warps=4),
-    ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK']
+    ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK']
 
 
 def get_rdna_autotune_configs():
@@ -242,7 +184,7 @@ def get_rdna_autotune_configs():
         # Fall-back config.
         triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1,
                       num_warps=2),
-    ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK']
+    ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK']
 
 
 def get_autotune_configs():
@@ -266,7 +208,7 @@ def get_autotune_configs():
             "MAX_SEQLENS_Q",
             "MAX_SEQLENS_K",
             "ACTUAL_BLOCK_DMODEL",
-            "VARLEN",
+            "IS_VARLEN",
             "HQ",
             "HK",
         ]
@@ -277,37 +219,47 @@ autotune_configs, autotune_keys = get_autotune_configs()
 @triton.autotune(
     configs=autotune_configs,
     key=autotune_keys,
-    # use_cuda_graph=True,
+    use_cuda_graph=True,
 )
 @triton.jit
-def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk,
+def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, # pylint: disable=unused-argument
+             SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk,
              stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn,
              stride_oz, stride_oh, stride_om, stride_on, stride_bz, stride_bh, stride_bm, stride_bn, stride_az, stride_ah, # pylint: disable=unused-argument
              stride_sz, stride_sh, stride_sm, stride_sn, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k,
-             dropout_p, philox_seed, philox_offset_base, scores, scores_scaled_shifted, exp_scores, alibi_slopes,  HQ: tl.constexpr,
+             dropout_p, philox_seed, philox_offset_base, sd_mask, dropout_mask, alibi_slopes, HQ: tl.constexpr,
              HK: tl.constexpr, ACTUAL_BLOCK_DMODEL: tl.constexpr, MAX_SEQLENS_Q: tl.constexpr,
-             MAX_SEQLENS_K: tl.constexpr, VARLEN: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr,
+             MAX_SEQLENS_K: tl.constexpr, IS_VARLEN: tl.constexpr, IS_INFERENCE: tl.constexpr,  IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr,
              BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, USE_BIAS: tl.constexpr,
              ENABLE_DROPOUT: tl.constexpr, RETURN_SCORES: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr):
+    # set params
+    ACCUMULATOR_TYPE = tl.float32
+
+    # compute offsets
     start_m = tl.program_id(0)
     off_h_q = tl.program_id(1)
     off_z = tl.program_id(2)
     offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n = tl.arange(0, BLOCK_N)
     offs_d = tl.arange(0, BLOCK_DMODEL)
-    if VARLEN:
+
+    # handle seqlen
+    if IS_VARLEN:
         cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
         cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
-        # print("cu_seqlens_q_start:", cu_seqlens_q_start)
-
         seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
-        # We have a one-size-fits-all grid in id(0). Some seqlens might be too
-        # small for all start_m so for those we return early.
+
+        # we have a one-size-fits-all grid in id(0). Some seqlens might be too small for all start_m so for those we return early.
         if start_m * BLOCK_M > seqlen_q:
             return
         cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
         cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
         seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
+    elif IS_INFERENCE:
+        cu_seqlens_q_start = 0
+        cu_seqlens_k_start = 0
+        seqlen_q = MAX_SEQLENS_Q
+        seqlen_k = tl.load(Cache_seqlens + off_z)
     else:
         cu_seqlens_q_start = 0
         cu_seqlens_k_start = 0
@@ -320,14 +272,14 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
     # inf written to LSE. We don't need to do any GEMMs in this case.
     # This block of code determines what N is, and if this WG is operating
     # on those M rows.
-    n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
+    n_blocks = tl.cdiv(seqlen_k, BLOCK_N)
     if IS_CAUSAL:
         # If seqlen_q == seqlen_k, the attn scores are a square matrix.
         # If seqlen_q != seqlen_k, attn scores are rectangular which means
         # the causal mask boundary is bottom right aligned, and ends at either
         # the top edge (seqlen_q < seqlen_k) or left edge.
         # This captures the decrease in n_blocks if we have a rectangular attn matrix
-        n_blocks_seqlen = cdiv_fn((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
+        n_blocks_seqlen = tl.cdiv((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
         # This is what adjusts the block_max for the current WG, only
         # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks
         n_blocks = min(n_blocks, n_blocks_seqlen)
@@ -345,7 +297,7 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
             l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m
             l_ptrs = l_offset + offs_m * stride_lse_m 
 
-            l = tl.full([BLOCK_M], value=0.0, dtype=tl.float32)
+            l = tl.full([BLOCK_M], value=0.0, dtype=ACCUMULATOR_TYPE)
 
             # mask_m_offsets = start_m + tl.arange(0, BLOCK_M)
             # lse_mask = mask_m_offsets < causal_start_idx
@@ -371,7 +323,7 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
         n_extra_tokens = BLOCK_N - seqlen_k
     elif seqlen_k % BLOCK_N:
         n_extra_tokens = seqlen_k % BLOCK_N
-    PADDED_HEAD: tl.constexpr = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
+    PADDED_HEAD: tl.constexpr = (ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL)
 
     # Compute pointers for all the tensors used in this kernel.
     q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm
@@ -394,28 +346,23 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
         alibi_slope = None
 
     if RETURN_SCORES:
-        scores_offset = scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm
-        score_ptrs = scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
-
-        scores_scaled_shifted_offset = scores_scaled_shifted + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm
-        scores_scaled_shifted_ptrs = scores_scaled_shifted_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
-
-        exp_scores_offset = exp_scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm
-        exp_scores_ptrs = exp_scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
+        sd_mask_offset = sd_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm
+        sd_mask_ptrs = sd_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
     else:
-        score_ptrs = None
-        scores_scaled_shifted_ptrs = None
-        exp_scores_ptrs = None
+        sd_mask_ptrs = None
 
     if ENABLE_DROPOUT:
-        off_hz = off_z * HQ + off_h_q
-        batch_philox_offset = philox_offset_base + off_hz * seqlen_q * seqlen_k
+        dropout_mask_offset = dropout_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm
+        dropout_mask_ptrs = dropout_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
+        batch_philox_offset = philox_offset_base + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm
+        philox_ptrs = batch_philox_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
     else:
-        batch_philox_offset = 0
+        dropout_mask_ptrs = None
+        philox_ptrs = 0
     # initialize pointer to m and l
-    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
-    l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
-    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=ACCUMULATOR_TYPE)
+    l_i = tl.full([BLOCK_M], 1.0, dtype=ACCUMULATOR_TYPE)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=ACCUMULATOR_TYPE)
     # Q is loaded once at the beginning and shared by all N blocks.
     q_ptrs_mask = offs_m[:, None] < seqlen_q
     if PADDED_HEAD:
@@ -442,16 +389,16 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
     # value because there is no masking. Similarly we do not need padding.
     if n_full_blocks > 0:
         block_max = (n_blocks - masked_blocks) * BLOCK_N
-        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn,
-                                        start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, batch_philox_offset,
-                                        exp_scores_ptrs,
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn,
+                                        start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, philox_ptrs,
+                                        sd_mask_ptrs, dropout_mask_ptrs,
                                         # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _
-                                        block_min, block_max, 0, 0, 0, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs,
+                                        block_min, block_max, 0, 0, 0, alibi_slope,
                                         # IS_CAUSAL, ....
                                         False, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n,
                                         # _, MASK_STEPS, ...
                                         PRE_LOAD_V, False, ENABLE_DROPOUT, PADDED_HEAD,
-                                        ACTUAL_BLOCK_DMODEL, SM_SCALE,  USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES)
+                                        ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE)
         block_min = block_max
         block_max = n_blocks * BLOCK_N
 
@@ -467,23 +414,25 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
         if USE_BIAS:
             bias_ptrs += n_full_blocks * BLOCK_N * stride_bn
         if RETURN_SCORES:
-            score_ptrs += n_full_blocks * BLOCK_N
-            scores_scaled_shifted_ptrs += n_full_blocks * BLOCK_N
-            exp_scores_ptrs += n_full_blocks * BLOCK_N
-        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn,
-                                        start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, batch_philox_offset,
-                                        exp_scores_ptrs, block_min, block_max, offs_n_causal, masked_blocks,
-                                        n_extra_tokens, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs,
+            sd_mask_ptrs += n_full_blocks * BLOCK_N * stride_sn
+        if ENABLE_DROPOUT:
+            dropout_mask_ptrs += n_full_blocks * BLOCK_N * stride_sn
+            philox_ptrs += n_full_blocks * BLOCK_N * stride_sn
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn,
+                                        start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, philox_ptrs,
+                                        sd_mask_ptrs, dropout_mask_ptrs, block_min, block_max, offs_n_causal, masked_blocks,
+                                        n_extra_tokens, alibi_slope,
                                         IS_CAUSAL, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n,
                                         # _, MASK_STEPS, ...
                                         PRE_LOAD_V, True, ENABLE_DROPOUT, PADDED_HEAD,
-                                        ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES)
+                                        ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE)
     # epilogue
     # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger.
     l_recip = 1 / l_i[:, None]
     acc = acc * l_recip
     if ENABLE_DROPOUT:
-        acc = acc / (1 - dropout_p)
+        dropout_scale = 1 / (1 - dropout_p)
+        acc = acc * dropout_scale
     # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,
     # then we have one block with a row of all NaNs which come from computing
     # softmax over a row of all -infs (-inf - inf = NaN). We check for that here
@@ -491,13 +440,12 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
     end_m_idx = (start_m + 1) * BLOCK_M
     start_m_idx = start_m * BLOCK_M
     causal_start_idx = seqlen_q - seqlen_k
-    acc = acc.to(Out.type.element_ty)
     if IS_CAUSAL:
         if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
             out_mask_boundary = tl.full((BLOCK_DMODEL, ), causal_start_idx, dtype=tl.int32)
             mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
             out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :]
-            z: tl.tensor = 0.0
+            z = 0.0
             acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
 
     # write back LSE(Log Sum Exponents), the log of the normalization constant
@@ -541,30 +489,43 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_
 
 
 def attention_prefill_forward_triton_impl(
-                                        q,
-                                        k,
-                                        v,
-                                        o,
-                                        sm_scale,
-                                        alibi_slopes,
-                                        causal,
-                                        bias,
-                                        dropout_p,
-                                        layout,
-                                        cu_seqlens_q,
-                                        cu_seqlens_k,
-                                        max_seqlens_q,
-                                        max_seqlens_k,
-                                        return_scores,
-                                        use_exp2):
-    # check if varlen
+                                        q: torch.Tensor,
+                                        k: torch.Tensor,
+                                        v: torch.Tensor,
+                                        o: torch.Tensor,
+                                        sm_scale: float,
+                                        alibi_slopes: Optional[torch.Tensor],
+                                        causal: bool,
+                                        bias: Optional[torch.Tensor],
+                                        layout: Literal["bshd", "bhsd", "thd"],
+                                        # varlen
+                                        cu_seqlens_q: Optional[torch.Tensor], 
+                                        cu_seqlens_k: Optional[torch.Tensor],
+                                        max_seqlens_q: int, 
+                                        max_seqlens_k: int,
+                                        # inference
+                                        cache_seqlens: Optional[Union[(int, torch.Tensor)]],
+                                        cache_batch_idx: Optional[torch.Tensor],
+                                        # dropout
+                                        dropout_p: float,
+                                        philox_seed: Optional[int],
+                                        philox_offset: Optional[int],
+                                        # misc
+                                        return_softmax: bool,
+                                        use_exp2: bool,
+):
+    # check flags
     is_varlen = layout == "thd"
+    use_alibi, (stride_az, stride_ah) = (True, alibi_slopes.stride()) if alibi_slopes is not None else (False, (0, 0))
+    is_inference = cache_seqlens is not None
+    if is_inference:
+        assert layout == "bshd", f"{layout} layout is not supported with inference. Use bshd layout"
 
     # NOTE: a large bias tensor leads to overflow during pointer arithmetic
-    if bias is not None:
-        assert bias.numel() < 2**31
+    if (bias is not None):
+        assert (bias.numel() < 2**31)
 
-    batch, nheads_q, nheads_k, head_size, seqlen_q, seqlen_k = get_shape_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k) # pylint: disable=unused-variable
+    batch, nheads_q, nheads_k, head_size, _, _ = get_shapes_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k)
     q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(q, k, v, o, layout)
 
     # Get closest power of 2 over or equal to 32.
@@ -573,59 +534,45 @@ def attention_prefill_forward_triton_impl(
     # kernel is padded - there is no padding in memory for any dims.
     padded_d_model = max(padded_d_model, 16)
 
-    grid = lambda META: (triton.cdiv(max_seqlens_q, META['BLOCK_M']), nheads_q, batch) # pylint: disable=unnecessary-lambda-assignment
+    grid = lambda META: (triton.cdiv(max_seqlens_q, META['BLOCK_M']), nheads_q, batch)
 
-    if return_scores:
-        scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
-                                        dtype=torch.float32)
-        scores_scaled_shifted = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
-                                        dtype=torch.float32)
-        scores_strides = (scores.stride(0), scores.stride(1), scores.stride(2), scores.stride(3))
-    else:
-        scores = None
-        scores_scaled_shifted = None
-        scores_strides = (0, 0 , 0 , 0)
-
-    # exp_scores is used to validate dropout behavior vs the PyTorch SDPA math backend reference.  We zero this out
+    # sd_mask is used to validate dropout behavior vs the PyTorch SDPA math backend reference.  We zero this out
     # to give a consistent starting point and then populate it with the output of softmax with the sign bit set according
     # to the dropout mask. The resulting return allows this mask to be fed into the reference implementation for testing
-    # only.  This return holds no useful output aside from debugging.
-    if return_scores:
-        exp_scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
-                                        dtype=torch.float32)
+    # only. This return holds no useful output aside from debugging.
+    use_dropout = (dropout_p > 0.0)
+    if use_dropout or return_softmax:
+        sd_mask = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
+                                dtype=torch.float32)
+        dropout_mask = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
+                                    dtype=torch.float32)
+        scores_strides = (sd_mask.stride(0), sd_mask.stride(1), sd_mask.stride(2), sd_mask.stride(3))
     else:
-        exp_scores = None
+        sd_mask = None
+        dropout_mask = None
+        scores_strides = (0, 0, 0, 0)
 
     # stores LSE the log of the normalization constant / sum of expoential score(unnormalzied probablities)
     if is_varlen:
-        softmax_lse = torch.empty((q.shape[0], nheads_q), device=q.device, dtype=torch.float32)
+        softmax_lse = torch.zeros((q.shape[0], nheads_q), device=q.device, dtype=torch.float32)
         stride_lse_m, stride_lse_h = softmax_lse.stride()
         stride_lse_z = 0
     else:
-        softmax_lse = torch.empty((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32)
+        softmax_lse = torch.zeros((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32)
         stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride()
 
-    # Seed the RNG so we get reproducible results for testing.
-    philox_seed = 0x1BF52
-    philox_offset = 0x1D4B42
-
     if bias is not None:
         bias_strides = (bias.stride(0), bias.stride(1),bias.stride(2),
                         bias.stride(3))
     else:
         bias_strides = (0, 0, 0, 0)
 
-    if alibi_slopes is not None:
-        alibi_strides = (alibi_slopes.stride(0), alibi_slopes.stride(1))
-    else:
-        alibi_strides = (0, 0)
-
-    attn_fwd[grid](q, k, v, bias, sm_scale, softmax_lse, o, *q_strides, *k_strides, *v_strides, *o_strides,
-                    *bias_strides, *alibi_strides, *scores_strides, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k,
-                    dropout_p=dropout_p, philox_seed=philox_seed, philox_offset_base=philox_offset, scores=scores,
-                    scores_scaled_shifted=scores_scaled_shifted, exp_scores=exp_scores, alibi_slopes=alibi_slopes,
+    attn_fwd[grid](q, k, v, bias, cache_seqlens, cache_batch_idx,
+                    sm_scale, softmax_lse, o, *q_strides, *k_strides, *v_strides, *o_strides,
+                    *bias_strides, stride_az, stride_ah, *scores_strides, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k,
+                    dropout_p=dropout_p, philox_seed=philox_seed, philox_offset_base=philox_offset, sd_mask=sd_mask, dropout_mask=dropout_mask, alibi_slopes=alibi_slopes, 
                     HQ=nheads_q, HK=nheads_k, ACTUAL_BLOCK_DMODEL=head_size, MAX_SEQLENS_Q=max_seqlens_q,
-                    MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, VARLEN=is_varlen,
+                    MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, IS_VARLEN=is_varlen, IS_INFERENCE=is_inference,
                     BLOCK_DMODEL=padded_d_model, USE_BIAS=False if bias is None else True,
-                    USE_ALIBI=False if alibi_slopes is None else True, ENABLE_DROPOUT=dropout_p
-                    > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_scores)
+                    USE_ALIBI=use_alibi, ENABLE_DROPOUT=dropout_p
+                    > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_softmax)

modules/flash_attn_triton_amd/interface_fa.py

@@ -1,22 +1,16 @@
 import torch
 from modules.flash_attn_triton_amd.fwd_prefill import attention_prefill_forward_triton_impl
-from modules.flash_attn_triton_amd.bwd_prefill import attention_prefill_backward_triton_impl
-from modules.flash_attn_triton_amd.fwd_decode import attention_decode_forward_triton_impl
-from modules.flash_attn_triton_amd.utils import MetaData, get_shape_from_layout
+from modules.flash_attn_triton_amd.utils import MetaData
 
 
-def fwd(q,
-    k,
-    v,
-    dropout_p,
-    softmax_scale,
-    causal,
+def fwd(q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        out: torch.Tensor,
+        dropout_p: float,
+        softmax_scale: float,
+        causal: bool
 ):
-    if dropout_p != 0.0:
-        raise ValueError("dropout is not supported on AMD's Triton Backend yet")
-
-    o = torch.empty_like(q)
-
     # Setup metadata
     metadata = MetaData(sm_scale=softmax_scale)
     metadata.max_seqlens_q = q.shape[1]
@@ -24,251 +18,35 @@ def fwd(q,
     metadata.layout = "bshd"
 
     if causal:
-        metadata.need_causal()
-
-    #if dropout_p > 0.0:
-    #    metadata.need_dropout(dropout_p, False)
-
-    # Check arguments
-    metadata.check_args(q, k, v, o)
-
-    attention_prefill_forward_triton_impl(
-                                        q,
-                                        k,
-                                        v,
-                                        o,
-                                        metadata.sm_scale,
-                                        metadata.alibi_slopes,
-                                        metadata.causal,
-                                        metadata.bias,
-                                        metadata.dropout_p,
-                                        metadata.layout,
-                                        metadata.cu_seqlens_q,
-                                        metadata.cu_seqlens_k,
-                                        metadata.max_seqlens_q,
-                                        metadata.max_seqlens_k,
-                                        metadata.return_scores,
-                                        metadata.use_exp2)
-
-    return o
-
-
-def bwd(
-    dout,
-    q,
-    k,
-    v,
-    out,
-    softmax_lse,
-    dq,
-    dk,
-    dv,
-    alibi_slopes,
-    dropout_p,
-    softmax_scale,
-    causal,
-    window_size_left, window_size_right, softcap, deterministic, gen_, rng_state, # pylint: disable=unused-argument
-):
-    if dropout_p != 0.0:
-        raise ValueError("dropout is not supported on AMD yet")
-
-    dq_triton, dk_triton, dv_triton, delta_triton, _, _ = attention_prefill_backward_triton_impl( # pylint: disable=unused-variable
-        dout,
-        q,
-        k,
-        v,
-        out,
-        softmax_lse,
-        dq,
-        dk,
-        dv,
-        softmax_scale,
-        alibi_slopes,
-        causal,
-        "bshd",
-        None,
-        None,
-        None,
-        None,
-        False,
-    )
-    delta = delta_triton
-
-    return dq, dk, dv, delta
-
-
-def varlen_fwd(
-    q,
-    k,
-    v,
-    o,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    seqused_k, leftpad_k, block_table_, # pylint: disable=unused-argument
-    alibi_slopes,\
-    max_seqlen_q,
-    max_seqlen_k,
-    dropout_p,
-    softmax_scale,
-    zero_tensors, # pylint: disable=unused-argument
-    causal,
-    window_size_left, window_size_right, softcap, # pylint: disable=unused-argument
-    return_softmax,
-    gen_ # pylint: disable=unused-argument
-):
-    if dropout_p != 0.0:
-        raise ValueError("dropout is not supported on AMD's Triton Backend yet")
-
-    if o is None:
-        o = torch.empty_like(q)
-
-    # Setup metadata
-    metadata = MetaData(sm_scale=softmax_scale)
-    if return_softmax:
-        metadata.return_scores = True
-    metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k)  # set layout to "thd" and other metdata
-
-    # get shapes
-    batch, nheads_q, nheads_k, head_size , seqlen_q, seqlen_k = get_shape_from_layout(q, k, metadata.layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) # pylint: disable=unused-variable
-
-    if causal:
-        metadata.need_causal()
-
-    if alibi_slopes is not None:
-        metadata.need_alibi(alibi_slopes, batch, nheads_q)
+        metadata.need_causal(True)
 
     if dropout_p > 0.0:
-        metadata.need_dropout(dropout_p, return_softmax)
+        metadata.need_dropout(dropout_p)
 
-    # Check arguments
-    metadata.check_args(q, k, v, o)
-    if o is None:
-        o = torch.empty_like(q, dtype=v.dtype)
+    # check arguments
+    metadata.check_args(q, k, v, out)
 
+    # call implementation
     attention_prefill_forward_triton_impl(
                                         q,
                                         k,
                                         v,
-                                        o,
+                                        out,
                                         metadata.sm_scale,
                                         metadata.alibi_slopes,
                                         metadata.causal,
-                                        metadata.bias,
-                                        metadata.dropout_p,
+                                        None,
                                         metadata.layout,
                                         metadata.cu_seqlens_q,
                                         metadata.cu_seqlens_k,
                                         metadata.max_seqlens_q,
                                         metadata.max_seqlens_k,
-                                        metadata.return_scores,
+                                        metadata.cache_seqlens,
+                                        metadata.cache_batch_idx,
+                                        metadata.dropout_p,
+                                        metadata.philox_seed,
+                                        metadata.philox_offset,
+                                        False,
                                         metadata.use_exp2)
 
-    return o
-
-
-def varlen_bwd(
-    dout,
-    q,
-    k,
-    v,
-    out,
-    softmax_lse,
-    dq,
-    dk,
-    dv,
-    cu_seqlens_q,
-    cu_seqlens_k,
-    alibi_slopes,
-    max_seqlen_q,
-    max_seqlen_k,
-    dropout_p,
-    softmax_scale,
-    zero_tensors, # pylint: disable=unused-argument
-    causal,
-    window_size_left, window_size_right, softcap, deterministic, gen_, rng_state, # pylint: disable=unused-argument
-):
-    if dropout_p != 0.0:
-        raise ValueError("dropout is not supported on AMD yet")
-
-    dq_triton, dk_triton, dv_triton, delta_triton, _, _ = attention_prefill_backward_triton_impl( # pylint: disable=unused-variable
-        dout,
-        q,
-        k,
-        v,
-        out,
-        softmax_lse,
-        dq,
-        dk,
-        dv,
-        softmax_scale,
-        alibi_slopes,
-        causal,
-        "thd",
-        cu_seqlens_q,
-        cu_seqlens_k,
-        max_seqlen_q,
-        max_seqlen_k,
-        False,
-    )
-    delta = delta_triton
-
-    return dq, dk, dv, delta
-
-
-def fwd_kvcache(
-    q,
-    k_cache,
-    v_cache,
-    k,
-    v,
-    cache_seqlens,
-    rotary_cos, rotary_sin, # pylint: disable=unused-argument
-    cache_batch_idx,
-    cache_leftpad, block_table, # pylint: disable=unused-argument
-    alibi_slopes,
-    out,
-    softmax_scale,
-    causal,
-    window_size_left, window_size_right, softcap, rotary_interleaved, num_splits, # pylint: disable=unused-argument
-):
-    if out is None:
-        out = torch.empty_like(q)
-
-    # fill metadata
-    metadata = MetaData(sm_scale=softmax_scale)
-    metadata.layout = "bshd"
-    metadata.max_seqlens_q = q.shape[1]
-    metadata.max_seqlens_k = k_cache.shape[1]
-    metadata.cache_seqlens = cache_seqlens
-    metadata.cache_batch_idx = cache_batch_idx
-
-    if k is not None and v is not None:
-        metadata.new_kv = True
-        metadata.seqlen_new = k.shape[1]
-        metadata.k_new = k
-        metadata.v_new = v
-
-    if causal:
-        metadata.need_causal()
-
-    if alibi_slopes is not None:
-        batch, _ , nheads_q, _= q.shape
-        metadata.need_alibi(alibi_slopes, batch, nheads_q)
-
-    # launch kernel
-    # TODO: pass output as an arg. Maybe we are copying output which is causing slow down
-    output, softmax_lse = attention_decode_forward_triton_impl(
-        q,
-        k_cache,
-        v_cache,
-        metadata.sm_scale,
-        metadata.causal,
-        metadata.alibi_slopes,
-        metadata.layout,
-        metadata.cache_seqlens,
-        metadata.cache_batch_idx,
-        metadata.new_kv,
-        metadata.k_new,
-        metadata.v_new,
-    )
-    return output, softmax_lse
+# varlen

modules/flash_attn_triton_amd/utils.py

@@ -1,33 +1,48 @@
-import os
+import csv
+import math
 import torch
+import os
+import random
+import functools
 import triton
+import triton.language as tl
+from typing import Literal, Optional, Union
 from modules.rocm import Agent, MicroArchitecture
 
 
 AUTOTUNE = os.environ.get('FLASH_ATTENTION_TRITON_AMD_AUTOTUNE', '0').lower() in ('1', 'true', 'yes')
+USE_REF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_REF', '0').lower() in ('1', 'true', 'yes')
 PERF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_PERF', '0').lower() in ('1', 'true', 'yes')
 
 
+# -------------------------------
+# Metadata
+# -------------------------------
 class MetaData():
-    cu_seqlens_q = None
-    cu_seqlens_k = None
-    max_seqlens_q = 0
-    max_seqlens_k = 0
-    bias = None
-    alibi_slopes = None
-    causal = False
+    cu_seqlens_q: Optional[torch.Tensor] = None
+    cu_seqlens_k: Optional[torch.Tensor] = None
+    max_seqlens_q: int = 0
+    max_seqlens_k: int = 0
+    bias: Optional[torch.Tensor] = None
+    alibi_slopes: Optional[torch.Tensor] = None
+    causal: bool = False
     num_contexts = 0
-    varlen = False
-    layout = None
-    cache_seqlens = None
+    varlen: bool = False
+    layout: Optional[Literal["bshd", "bhsd", "thd"]] = None
+    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None
     cache_batch_idx = None
-    new_kv = False
-    seqlen_new = None
-    k_new = None
-    v_new = None
-    dropout_p, return_scores= 0.0, False
+    packing: Optional[bool] = None
+    return_scores: bool = False
+    dropout_p: float = 0.0
+    philox_seed: Optional[int] = None
+    philox_offset : Optional[int]= None # if dropout_p > 0.0 seed the RNG so we get reproducible results for testing.
     # NOTE: scale sm_scale by log_2(e) and use 2^x in the loop as we do not have native e^x support in HW.
-    use_exp2 = False
+    use_exp2: bool = False
+    rotary_sin: Optional[torch.Tensor] = None
+    rotary_cos: Optional[torch.Tensor] = None
+    rotary_interleaved: bool = False
+    rotary_conjunction: bool = False
+    
 
     def __repr__(self) -> str:
         return (f"MetaData(\n"
@@ -44,10 +59,6 @@ class MetaData():
                 f"  layout={self.layout},\n"
                 f"  cache_seqlens={self.cache_seqlens},\n"
                 f"  cache_batch_idx={self.cache_batch_idx},\n"
-                f"  new_kv={self.new_kv},\n"
-                f"  seqlen_new={self.seqlen_new},\n"
-                f"  k_new={self.k_new},\n"
-                f"  v_new={self.v_new},\n"
                 f"  dropout_p={self.dropout_p},\n"
                 f"  return_scores={self.return_scores}\n"
                 f")")
@@ -55,20 +66,19 @@ class MetaData():
     def __init__(self, sm_scale=1.0):
         self.sm_scale = sm_scale
 
-    def set_varlen_params(self, cu_seqlens_q, cu_seqlens_k):
+    def set_varlen_params(self, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k):
         self.varlen = True
         self.layout = 'thd'
         self.cu_seqlens_q = cu_seqlens_q
         self.cu_seqlens_k = cu_seqlens_k
+        self.max_seqlens_q = max_seqlen_q
+        self.max_seqlens_k = max_seqlen_k
+
         # Without "varlen", there should still be one sequence.
         assert len(cu_seqlens_q) >= 2
         assert len(cu_seqlens_q) == len(cu_seqlens_k)
-        self.num_contexts = len(cu_seqlens_q) - 1
-        for i in range(0, self.num_contexts):
-            self.max_seqlens_q = max(cu_seqlens_q[i + 1].item() - cu_seqlens_q[i].item(), self.max_seqlens_q)
-            self.max_seqlens_k = max(cu_seqlens_k[i + 1].item() - cu_seqlens_k[i].item(), self.max_seqlens_k)
 
-    def need_bias(self, bias, batch, nheads, seqlen_q, seqlen_k): # pylint: disable=unused-argument
+    def need_bias(self, bias, batch, nheads, seqlen_q, seqlen_k):
         assert bias.is_cuda
         assert bias.dim() == 4
         assert bias.shape[0] == 1
@@ -82,17 +92,25 @@ class MetaData():
         assert alibi_slopes.shape[1] == nheads
         self.alibi_slopes = alibi_slopes
 
-    def need_causal(self):
-        self.causal = True
+    def need_causal(self, causal):
+        self.causal = causal
 
-    def need_dropout(self, dropout_p, return_scores):
-        self.dropout_p = dropout_p
-        self.return_scores = return_scores
+    def need_rotary(self, sin, cos, rotary_interleaved, rotary_conjunction=False):
+        self.rotary_sin = sin
+        self.rotary_cos = cos
+        self.rotary_interleaved = rotary_interleaved
+        self.rotary_conjunction = rotary_conjunction
+
+    def need_dropout(self, dropout_p, return_scores = True):
+        if dropout_p > 0.0:
+            self.dropout_p = dropout_p
+            self.return_scores = return_scores
+            self.philox_seed, self.philox_offset = 0x1BF58, 0x1D4B49
 
     def check_args(self, q, k, v, o):
         assert q.dim() == k.dim() and q.dim() == v.dim()
 
-        batch, nheads_q, nheads_k, head_size, _, _ = get_shape_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k) # pylint: disable=unused-variable
+        batch, nheads_q, nheads_k, head_size, _, _ = get_shapes_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k)
         if self.varlen:
             assert q.dim() == 3
             assert self.cu_seqlens_q is not None
@@ -100,8 +118,6 @@ class MetaData():
             assert len(self.cu_seqlens_q) == len(self.cu_seqlens_k)
             # TODO: Remove once bias is supported with varlen
             assert self.bias is None
-            # TODO:Remove once dropout is supported with varlen
-            assert self.dropout_p == 0.0
             # assert not self.return_scores
         else:
             assert q.dim() == 4
@@ -111,138 +127,286 @@ class MetaData():
         assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
         # TODO: Change assert if we support qkl f8 and v f16
         assert q.dtype == k.dtype and q.dtype == v.dtype
-        assert head_size <= 256
         assert o.shape == q.shape
         assert (nheads_q % nheads_k) == 0
         assert self.layout is not None
         assert self.layout == 'thd' or not self.varlen
 
-def input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, layout, device="cuda", DEBUG_INPUT=False):
+# -------------------------------
+# Input Helper
+# -------------------------------
+def random_seqlens_composition(SEQ_LEN, BATCH):
+    # generate a random composition of N into Z positive parts.
+    idx = torch.randperm(SEQ_LEN - 1)[: BATCH - 1] + 1
+    idx, _ = torch.sort(idx)
+    breakpoints = torch.cat([
+        torch.tensor([0], dtype=torch.long),
+        idx,
+        torch.tensor([SEQ_LEN], dtype=torch.long),
+    ])
+    seqlens = (breakpoints[1:] - breakpoints[:-1]).to(torch.int32)
+    return seqlens
+
+def generate_varlen_tensor(
+    total_seqlen: int,
+    num_heads: int,
+    head_size: int,
+    batch_size: Optional[int] = None,
+    equal_seqlens: bool = False,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.float32,
+    DEBUG_INPUT: bool = False
+):
+    # get valid batch_size
+    if batch_size is None:
+        valid_batch_sizes = [bs for bs in [1, 2, 4, 8, 16, 32, 64] if bs <= total_seqlen]
+        batch_size = random.choice(valid_batch_sizes)
+    
+    # get seqlens
+    if equal_seqlens:
+        seqlens = torch.full(
+        (batch_size,),
+        total_seqlen // batch_size,
+        dtype=torch.int32,
+        device=device
+        )
+        seqlens[-1] += total_seqlen % batch_size
+    else:
+        seqlens = random_seqlens_composition(total_seqlen, batch_size).to(device=device)
+
+    # create cumulative sequence lengths
+    cu_seqlens = torch.cat([torch.tensor([0], dtype=torch.int32, device=device), seqlens.cumsum(dim=0)]).to(torch.int32).to(device=device)
+    max_seqlen = torch.max(seqlens).to(torch.int32).item()
+
+    # create varlen tensor
+    if DEBUG_INPUT:
+        x = torch.zeros(total_seqlen, num_heads, head_size, dtype=dtype, device=device)
+        for i in range(batch_size):
+            start = cu_seqlens[i].item()
+            end   = cu_seqlens[i+1].item()
+            length  = end - start
+
+            x[start:end, :, :] = (
+                torch.arange(length, dtype=dtype, device=device)
+                .view(length, 1, 1)
+                .expand(length, num_heads, head_size)
+            )
+    else:
+        x = torch.randn((total_seqlen, num_heads, head_size), dtype=dtype, device=device)
+
+    x.requires_grad_()
+    return x, cu_seqlens, max_seqlen
+
+def generate_bshd_tensor(BATCH, SEQ_LEN, NUM_HEADS, D_HEAD, dtype, device="cuda", DEBUG_INPUT=False):
+    # gen tensor
+    tensor_shape = (BATCH, SEQ_LEN, NUM_HEADS, D_HEAD)
+    if DEBUG_INPUT:
+        x = torch.arange(SEQ_LEN, dtype=dtype, device=device).view(1, SEQ_LEN, 1, 1).expand(*tensor_shape).contiguous()
+    else:
+        x = torch.randn(tensor_shape, dtype=dtype, device=device)
+
+    x.requires_grad_()
+    return x
+
+def generate_bhsd_tensor(BATCH, NUM_HEADS, SEQ_LEN, D_HEAD, dtype, device="cuda", DEBUG_INPUT=False):
+    # gen tensor
+    tensor_shape = (BATCH, NUM_HEADS, SEQ_LEN, D_HEAD)
+    if DEBUG_INPUT:
+        x = torch.arange(SEQ_LEN, dtype=dtype, device=device).view(1, 1, SEQ_LEN, 1).expand(*tensor_shape).contiguous()
+    else:
+        x = torch.randn(tensor_shape, dtype=dtype, device=device)
+
+    x.requires_grad_()
+    return x
+
+def input_helper(
+    BATCH: int,
+    HQ: int,
+    HK: int,
+    N_CTX_Q: int,
+    N_CTX_K: int,
+    D_HEAD: int,
+    CAUSAL: bool,
+    DROPOUT_P: float,
+    dtype: torch.dtype,
+    layout: Literal["bshd", "bhsd", "thd"],
+    packing: Optional[Literal["kv", "qkv"]] = None,
+    device: Literal["cpu", "cuda"] = "cuda",
+    DEBUG_INPUT: bool = False,
+):
     torch.manual_seed(20)
 
-    # Initialize q, k, v
-    if layout == 'bhsd':
-        q_tensor_shape = (Z, HQ, N_CTX_Q, D_HEAD)
-        k_tensor_shape = (Z, HK, N_CTX_K, D_HEAD)
-    elif layout == 'bshd':
-        q_tensor_shape = (Z, N_CTX_Q, HQ, D_HEAD)
-        k_tensor_shape = (Z, N_CTX_K, HK, D_HEAD)
+    if layout == "thd":
+        # set params
+        TOTAL_SEQLENS_Q = BATCH * N_CTX_Q
+        TOTAL_SEQLENS_K = BATCH * N_CTX_K
+        equal_seqlens=False
+        
+        # gen tensors
+        # TODO: the gen functions should maybe have different gen modes like random, ones, increasing seqlen
+        q, cu_seqlens_q, max_seqlen_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT)
+        k, cu_seqlens_k, max_seqlen_k = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT)
+        v, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT)
+        do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q)
+        
+        # setup metadata
+        if DEBUG_INPUT:
+            sm_scale = 1
+        else:
+            sm_scale = D_HEAD**-0.5
+        metadata = MetaData(sm_scale=sm_scale)
+        metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k)
+        metadata.need_causal(CAUSAL)
+        metadata.need_dropout(DROPOUT_P)
+    elif layout == 'bshd' or layout == "bhsd":
+        # gen tensors
+        if layout == "bshd":
+            q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT)
+            k = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT)
+            v = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT)
+            do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q)
+        elif layout == "bhsd":
+            q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT)
+            k = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT)
+            v = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT)
+            do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q)
+
+        # setup metadata
+        if DEBUG_INPUT:
+            sm_scale = 1
+        else:
+            sm_scale = D_HEAD**-0.5
+        metadata = MetaData(sm_scale=sm_scale)
+        metadata.max_seqlens_q = N_CTX_Q
+        metadata.max_seqlens_k = N_CTX_K
+        metadata.layout = layout
+        metadata.need_causal(CAUSAL)
+        metadata.need_dropout(DROPOUT_P)
     else:
-        assert False, f'Got unsupported tensor layout: {layout}'
+        raise ValueError(f"Unknown layout: {layout}")
 
-    q = None
-    k = None
-    v = None
-
-    if DEBUG_INPUT:
-        if layout == "bhsd":
-            q = torch.arange(N_CTX_Q, dtype=dtype, device=device).view(1, 1, N_CTX_Q, 1).expand(*q_tensor_shape).contiguous().requires_grad_()
-            k = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, 1, N_CTX_K, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
-            v = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, 1, N_CTX_K, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
+    # deal with packing
+    if packing is None:
+        return q, k, v, do, metadata
+    elif packing == "kv":
+        # pack k and v
+        if layout in ["bhsd", "thd"]:
+            kv = torch.stack([k, v], dim=1)
         elif layout == "bshd":
-            q = torch.arange(N_CTX_Q, dtype=dtype, device=device).view(1, N_CTX_Q, 1, 1).expand(*q_tensor_shape).contiguous().requires_grad_()
-            k = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, N_CTX_K, 1, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
-            v = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, N_CTX_K, 1, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
+            kv = torch.stack([k, v], dim=2)
+        else:
+            raise ValueError(f"Unknown layout: {layout}")
+
+        return q, kv, do, metadata
+    elif packing == "qkv":
+        # qkv packing - requires same sequence length for q and k
+        assert N_CTX_Q == N_CTX_K, "For QKV packing, Q and K must have same sequence length"
+        assert HQ == HK, "For QKV packing, Q and K must have same number of heads"
+
+        # pack q, k, and v
+        if layout in ["bhsd", "thd"]:
+            qkv = torch.stack([q, k, v], dim=1)
+        elif layout == "bshd":
+            qkv = torch.stack([q, k, v], dim=2)
+        else:
+            raise ValueError(f"Unknown layout: {layout}")
+
+        return qkv, do, metadata
     else:
-        q = torch.randn(q_tensor_shape, dtype=dtype, device=device, requires_grad=True)
-        k = torch.randn(k_tensor_shape, dtype=dtype, device=device, requires_grad=True)
-        v = torch.randn(k_tensor_shape, dtype=dtype, device=device, requires_grad=True)
+        assert False, f"Unsupported packing mode: {packing}"
 
-    if DEBUG_INPUT:
-        sm_scale = 1
+# -------------------------------
+# Alibi
+# -------------------------------
+@triton.jit
+def compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False):
+    # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix
+    # for casual mask we want something like this where (1 is kept and 0 is masked)
+    # seqlen_q = 2 and seqlen_k = 5
+    #   1 1 1 1 0
+    #   1 1 1 1 1
+    # seqlen_q = 5 and seqlen_k = 2
+    #        0 0
+    #        0 0
+    #        0 0
+    #        1 0
+    #        1 1
+    # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal
+    # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5, offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False
+    # 1. offs_m[:,None] = [[0],
+    #                       [1],
+    # 2. offs_m[:,None] + seqlen_k = [[5],
+    #                                  [6],
+    # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3],
+    #                                             [4],
+    # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] = [[3], - [[0, 1, 2, 3, 4]] =  [[ 3, 2, 1, 0,-1],
+    #                                                            [4],                           [ 4, 3, 2, 1, 0]]
+    # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1],
+    #                                                     [ -4, -3, -2, -1, 0]],
+    relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :]
+    alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
+    if transpose:
+        return alibi_block.T
     else:
-        sm_scale = D_HEAD**-0.5
-    input_metadata = MetaData(sm_scale=sm_scale)
-    input_metadata.max_seqlens_q = N_CTX_Q
-    input_metadata.max_seqlens_k = N_CTX_K
-    input_metadata.layout = layout
-    return q, k, v, input_metadata
+        return alibi_block
 
-
-def varlen_input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, device="cuda", equal_seqlens=False, DEBUG_INPUT=False):
-    torch.manual_seed(20)
-
-    # Random or equal sequence lengths based on 'equal_seqlens' flag
-    if not equal_seqlens:
-        max_seqlens_q = N_CTX_Q // Z
-        max_seqlens_k = N_CTX_K // Z
-        seqlens_q = torch.randint(1, max_seqlens_q + 1, (Z,), dtype=torch.int32)
-        seqlens_k = torch.randint(1, max_seqlens_k + 1, (Z,), dtype=torch.int32)
-    else:
-        seqlens_q = torch.full((Z,), N_CTX_Q // Z, dtype=torch.int32)
-        seqlens_k = torch.full((Z,), N_CTX_K // Z, dtype=torch.int32)
-
-    # Calculate cumulative sequence lengths
-    cu_seqlens_q = torch.cat([torch.tensor([0], dtype=torch.int32), seqlens_q.cumsum(dim=0)])
-    cu_seqlens_k = torch.cat([torch.tensor([0], dtype=torch.int32), seqlens_k.cumsum(dim=0)])
-    cu_seqlens_q = cu_seqlens_q.to(device=device).to(torch.int32)
-    cu_seqlens_k = cu_seqlens_k.to(device=device).to(torch.int32)
-
-    # Total lengths
-    total_q = cu_seqlens_q[-1].item()
-    total_k = cu_seqlens_k[-1].item()
-
-    if DEBUG_INPUT:
-        # Initialize q, k, v with deterministic values
-        q = torch.arange(total_q, dtype=dtype, device=device).view(total_q, 1, 1)
-        q = q.expand(total_q, HQ, D_HEAD).contiguous().requires_grad_()
-        k = torch.arange(total_k, dtype=dtype, device=device).view(total_k, 1, 1)
-        k = k.expand(total_k, HK, D_HEAD).contiguous().requires_grad_()
-        v = torch.arange(total_k, dtype=dtype, device=device).view(total_k, 1, 1)
-        v = v.expand(total_k, HK, D_HEAD).contiguous().requires_grad_()
-        sm_scale = 1
-    else:
-        # Initialize q, k, v with random values
-        q = torch.randn((total_q, HQ, D_HEAD), dtype=dtype, device=device).requires_grad_()
-        k = torch.randn((total_k, HK, D_HEAD), dtype=dtype, device=device).requires_grad_()
-        v = torch.randn((total_k, HK, D_HEAD), dtype=dtype, device=device).requires_grad_()
-        sm_scale = D_HEAD ** -0.5
-
-    input_metadata = MetaData(sm_scale=sm_scale)
-    input_metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k)
-    return q, k, v, input_metadata
-
-
-def get_shape_from_layout(q, k, layout, cu_seqlens_q = None, cu_seqlens_k = None, max_seqlen_q=None, max_seqlen_k=None):
+# -------------------------------
+# Misc
+# -------------------------------
+def get_shape_from_layout(
+    x: torch.Tensor,
+    layout: Literal["bshd", "bhsd", "thd"],
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+) -> tuple[int, int, int, int]:
     if layout == 'bhsd':
-        batch_q, nheads_q, max_seqlen_q, head_size_q = q.shape
-        batch_k, nheads_k, max_seqlen_k, head_size_k = k.shape
+        batch, num_heads, max_seqlen_final, head_dim = x.shape
     elif layout == 'bshd':
-        batch_q, max_seqlen_q, nheads_q, head_size_q = q.shape
-        batch_k, max_seqlen_k, nheads_k, head_size_k = k.shape
+        batch, max_seqlen_final, num_heads, head_dim = x.shape
     elif  layout == 'thd':
-        batch_q, max_seqlen_q, nheads_q, head_size_q = len(cu_seqlens_q) - 1, max_seqlen_q, q.shape[1], q.shape[2] # pylint: disable=self-assigning-variable
-        batch_k, max_seqlen_k, nheads_k, head_size_k = len(cu_seqlens_k) - 1, max_seqlen_k, k.shape[1], k.shape[2] # pylint: disable=self-assigning-variable
+        total_seqlen, num_heads, head_dim = x.shape
+        if cu_seqlens is None:
+            raise ValueError("cu_seqlens must be provided for varlen (thd) layout") 
+        if max_seqlen is None:
+            raise ValueError("max_seqlen must be provided for varlen (thd) layout")
+        
+        batch, max_seqlen_final, num_heads, head_dim = len(cu_seqlens) - 1, max_seqlen, num_heads, head_dim
     else:
         assert False, "Got unsupported layout."
 
+    return batch, max_seqlen_final, num_heads, head_dim
+
+
+def get_shapes_from_layout(q, k, layout, cu_seqlens_q = None, cu_seqlens_k = None, max_seqlen_q=None, max_seqlen_k=None):
+    batch_q, seqlen_q, nheads_q, head_size_q = get_shape_from_layout(q, layout, cu_seqlens_q, max_seqlen_q)
+    batch_k, seqlen_k, nheads_k, head_size_k = get_shape_from_layout(k, layout, cu_seqlens_k, max_seqlen_k)
+    
     # assert
     assert batch_q == batch_k
     assert head_size_q == head_size_k
 
-    return batch_q, nheads_q, nheads_k, head_size_q, max_seqlen_q, max_seqlen_k
+    return batch_q, nheads_q, nheads_k, head_size_q, seqlen_q, seqlen_k
 
-
-def get_strides_from_layout(q, k, v, o, layout):
+def get_stride_from_layout(x: torch.Tensor, layout:Literal["bshd", "bhsd", "thd"]):
     if layout == 'thd':
-        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
-        k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
-        v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
-        o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+        strides = (0, x.stride(1), x.stride(0), x.stride(2))  
     elif layout == 'bhsd':
-        q_strides = (q.stride(0), q.stride(1), q.stride(2), q.stride(3))
-        k_strides = (k.stride(0), k.stride(1), k.stride(2), k.stride(3))
-        v_strides = (v.stride(0), v.stride(1), v.stride(2), v.stride(3))
-        o_strides = (o.stride(0), o.stride(1), o.stride(2), o.stride(3))
+        strides = (x.stride(0), x.stride(1), x.stride(2), x.stride(3))
     elif layout == 'bshd':
-        q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
-        k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
-        v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
-        o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+        strides = (x.stride(0), x.stride(2), x.stride(1), x.stride(3))
     else:
         assert False, 'Got unsupported layout.'
-    return q_strides, k_strides, v_strides, o_strides
+    return strides
 
+def get_shape_and_strides_from_layout(x: torch.Tensor, layout: Literal["bshd", "bhsd", "thd"], cu_seqlens: Optional[torch.Tensor] = None, max_seqlen: Optional[int] = None):
+    return get_shape_from_layout(x, layout, cu_seqlens, max_seqlen), get_stride_from_layout(x, layout)
+
+def get_strides_from_layout(q, k, v, o, layout):
+    q_strides = get_stride_from_layout(q, layout)
+    k_strides = get_stride_from_layout(k, layout)
+    v_strides = get_stride_from_layout(v, layout)
+    o_strides = get_stride_from_layout(o, layout)
+    return q_strides, k_strides, v_strides, o_strides
 
 def get_padded_headsize(size):
     # Get closest power of 2 over or equal to 32.
@@ -252,24 +416,79 @@ def get_padded_headsize(size):
     padded_d_model = max(padded_d_model, 16)
     return padded_d_model
 
+def compute_alibi_tensor_ref(alibi_slopes, seqlen_q, seqlen_k):
+    q_idx = torch.arange(seqlen_q, dtype=torch.int32, device="cuda").unsqueeze(-1)  # (N_CTX_Q, 1)
+    k_idx = torch.arange(seqlen_k, dtype=torch.int32, device="cuda").unsqueeze(0)  # (1, N_CTX_K)
+    relative_pos = torch.abs(q_idx + seqlen_k - seqlen_q - k_idx)  # (N_CTX_Q, N_CTX_K)
+    return -1 * alibi_slopes.unsqueeze(-1).unsqueeze(-1) * relative_pos  # (Z, H, N_CTX_Q, N_CTX_K)
 
-def _strides(x: torch.Tensor, *stride_names: str):
-    if x is None:
-        return {f"stride_{s}": 0 for i, s in enumerate(stride_names)}
+# -------------------------------
+# Dropouts
+# -------------------------------
+def create_dropout_mask(dropout_p, shape, seed):
+    device = "cuda"
+    rand_vals = torch.rand(shape, generator=torch.Generator(device=device).manual_seed(seed), device=device, dtype=torch.float32)
+    return rand_vals > dropout_p
 
-    assert x.ndim == len(stride_names)
-    return {f"stride_{s}": x.stride(i) for i, s in enumerate(stride_names)}
+def create_dropout_mask_varlen(dropout_p, batch, nheads_q, cu_seqlens_q, cu_seqlens_k, philox_seed):
+    device = "cuda"
+    qlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1])
+    klens = (cu_seqlens_k[1:] - cu_seqlens_k[:-1])
+    max_qlen = qlens.max()
+    max_klen = klens.max()
+    dropout_mask = torch.zeros((batch, nheads_q, max_qlen, max_klen), device=device)
+    for b in range(batch):
+        qlen = qlens[b]
+        klen = klens[b]
+        rand_vals = torch.rand((nheads_q, qlen, klen), generator=torch.Generator(device=device).manual_seed(philox_seed), device=device, dtype=torch.float32)
+        submask = rand_vals > dropout_p
+        dropout_mask[b, :, :qlen, :klen] = submask
 
+    return dropout_mask
 
-def get_input_shapes():
-    cases = [(max(1, 2**(16 - i)), 1, 2**i, 16, 1, 128)
-             for i in range(8, 18)] + [(max(1, 2**(16 - i)), 1, 2**i, 16, 2, 128) for i in range(8, 18)]
-    return cases
+def write_dropout_mask(x, tensor_name = "tensor"):
+    batch, head, seqlen_m, seqlen_n = x.shape
+    x = x.tolist()
 
+    with open(f'{tensor_name}.csv', 'w') as f:
+        writer = csv.writer(f)
+        for b in range(batch):
+            for h in range(head):
+                dropout_mask = x[b][h]
+                if True:
+                    BLOCK_M = 64
+                    BLOCK_N = 64
+                
+                    # Calculate number of blocks in each dimension
+                    m_blocks = math.ceil(seqlen_m / BLOCK_M)
+                    n_blocks = math.ceil(seqlen_n / BLOCK_N)
+                    
+                    # Process each block
+                    for m_block in range(m_blocks):
+                        # Calculate row range for current block
+                        row_start = m_block * BLOCK_M
+                        row_end = min(row_start + BLOCK_M, seqlen_m)
+                        
+                        for n_block in range(n_blocks):
+                            # Calculate column range for current block
+                            col_start = n_block * BLOCK_N
+                            col_end = min(col_start + BLOCK_N, seqlen_n)
+                            
+                            # Extract and write the current block
+                            for row_idx in range(row_start, row_end):
+                                row_data = dropout_mask[row_idx][col_start:col_end]
+                                writer.writerow(row_data)
+                else:
+                    writer.writerows(dropout_mask)
 
+# -------------------------------
+# Runtime info
+# -------------------------------
+@functools.cache
 def is_cdna():
     return Agent(triton.runtime.driver.active.get_current_target().arch).arch == MicroArchitecture.CDNA
 
 
+@functools.cache
 def is_rdna():
     return Agent(triton.runtime.driver.active.get_current_target().arch).arch == MicroArchitecture.RDNA

modules/infotext.py

@@ -68,7 +68,6 @@ def parse(infotext):
     debug(f'Raw: {infotext}')
 
     remaining = infotext.replace('\nSteps:', ' Steps:')
-    # TODO infotext: handle using regex instead
     params = [' steps:', ' seed:', ' width:', ' height:', ' sampler:', ' size:', ' cfg scale:'] # first param is one of those
     params += ['\nsteps:', '\nseed:', '\nwidth:', '\nheight:', '\nsampler:', '\nsize:', '\ncfg scale:']
     params += ['.steps:', '.seed:', '.width:', '.height:', '.sampler:', '.size:', '.cfg scale:']

modules/intel/ipex/diffusers.py

@@ -1,6 +1,7 @@
 from functools import wraps
 import torch
 import diffusers # pylint: disable=import-error
+from diffusers.utils import torch_utils # pylint: disable=import-error, unused-import # noqa: F401
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
@@ -61,13 +62,11 @@ def hidream_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
 
 
 def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False):
-    # get around lazy imports
-    from diffusers.utils import torch_utils # pylint: disable=import-error, unused-import
     diffusers.utils.torch_utils.fourier_filter = fourier_filter
     if not device_supports_fp64:
         # get around lazy imports
-        from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import
-        from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import
+        from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import # noqa: F401
+        from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import # noqa: F401
         diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed
         diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed
         diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed

modules/interrogate/openclip.py

@@ -350,7 +350,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod
         files += [f.name for f in batch_folder]
     if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
         from modules.files_cache import list_files
-        files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp'], recursive=recursive))
+        files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive))
     if len(files) == 0:
         shared.log.warning('Interrogate batch: type=clip no images')
         return ''

modules/interrogate/vqa.py

@@ -27,6 +27,7 @@ vlm_models = {
     "Google Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB
     "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224",
     "Alibaba Qwen VL2 2B": "Qwen/Qwen2-VL-2B-Instruct",
+    "Alibaba Qwen 2.5 Omni 3B": "Qwen/Qwen2.5-Omni-3B",
     "Huggingface Smol VL2 0.5B": "HuggingFaceTB/SmolVLM-500M-Instruct",
     "Huggingface Smol VL2 2B": "HuggingFaceTB/SmolVLM-Instruct",
     "Salesforce BLIP Base": "Salesforce/blip-vqa-base", # 1.5GB
@@ -67,6 +68,8 @@ vlm_prompts = [
 
 
 def b64(image):
+    if image is None:
+        return ''
     with io.BytesIO() as stream:
         image.save(stream, 'JPEG')
         values = stream.getvalue()
@@ -125,7 +128,7 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
-    system_prompt = system_prompt or shared.opts.vlm_system
+    system_prompt = system_prompt or shared.opts.interrogate_vlm_system
     conversation = [
         {
             "role": "system",
@@ -169,19 +172,20 @@ def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: st
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
-    system_prompt = system_prompt or shared.opts.vlm_system
+    system_prompt = system_prompt or shared.opts.interrogate_vlm_system
+
+    system_content = []
+    if system_prompt is not None and len(system_prompt) > 4:
+        system_content.append({"type": "text", "text": system_prompt})
+
+    user_content = []
+    if image is not None:
+        user_content.append({"type": "image", "image": b64(image)})
+    if question is not None and len(question) > 4:
+        user_content.append({"type": "text", "text": question})
     conversation = [
-        {
-            "role": "system",
-            "content": [{"type": "text", "text": system_prompt}]
-        },
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": b64(image)},
-                {"type": "text", "text": question}
-            ]
-        }
+        { "role": "system", "content": system_content},
+        { "role": "user", "content": user_content },
     ]
     inputs = processor.apply_chat_template(
         conversation,
@@ -294,7 +298,7 @@ def smol(question: str, image: Image.Image, repo: str = None, system_prompt: str
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
-    system_prompt = system_prompt or shared.opts.vlm_system
+    system_prompt = system_prompt or shared.opts.interrogate_vlm_system
     conversation = [
         {
             "role": "system",
@@ -632,7 +636,7 @@ def batch(model_name, system_prompt, batch_files, batch_folder, batch_str, quest
         files += [f.name for f in batch_folder]
     if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
         from modules.files_cache import list_files
-        files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp'], recursive=recursive))
+        files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive))
     if len(files) == 0:
         shared.log.warning('Interrogate batch: type=vlm no images')
         return ''

modules/model_flex.py

@@ -82,8 +82,8 @@ def load_flex(checkpoint_info, diffusers_load_config={}):
     )
     sd_hijack_te.init_hijack(pipe)
     diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline
-    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["fluxcfgzero"] = Flex2Pipeline
-    diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["fluxcfgzero"] = Flex2Pipeline
+    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline
+    diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flex2"] = Flex2Pipeline
 
     del text_encoder_2
     del transformer

modules/model_hidream.py

@@ -42,6 +42,8 @@ def load_transformer(repo_id, diffusers_load_config={}):
 
 
 def load_text_encoders(repo_id, diffusers_load_config={}):
+    if repo_id == 'HiDream-ai/HiDream-E1-Full':
+        repo_id = 'HiDream-ai/HiDream-I1-Full' # use I1 for t5 and llm
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     text_encoder_3 = transformers.T5EncoderModel.from_pretrained(
@@ -92,7 +94,19 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
     shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
 
-    pipe = diffusers.HiDreamImagePipeline.from_pretrained(
+    if 'I1' in repo_id:
+        cls = diffusers.HiDreamImagePipeline
+    elif 'E1' in repo_id:
+        from modules.hidream.pipeline_hidream_image_editing import HiDreamImageEditingPipeline
+        cls = HiDreamImageEditingPipeline
+        diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["hidream-e1"] = diffusers.HiDreamImagePipeline
+        diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["hidream-e1"] = HiDreamImageEditingPipeline
+        diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["hidream-e1"] = HiDreamImageEditingPipeline
+    else:
+        shared.log.error(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" not recognized')
+        return False
+
+    pipe = cls.from_pretrained(
         repo_id,
         transformer=transformer,
         text_encoder_3=text_encoder_3,
@@ -101,6 +115,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         cache_dir=shared.opts.diffusers_dir,
         **load_args,
     )
+
     sd_hijack_te.init_hijack(pipe)
     del text_encoder_3
     del text_encoder_4

modules/model_quant.py

@@ -207,7 +207,7 @@ def load_bnb(msg='', silent=False):
     if not installed('bitsandbytes'):
         if devices.backend == 'cuda':
             # forcing a version will uninstall the multi-backend-refactor branch of bnb
-            install('bitsandbytes==0.45.1', quiet=True)
+            install('bitsandbytes==0.45.5', quiet=True)
             log.warning('Quantization: bitsandbytes installed please restart')
     try:
         import bitsandbytes
@@ -263,6 +263,8 @@ def load_nncf(msg='', silent=False):
     if not installed('nncf'):
         install('nncf==2.16.0', quiet=True)
         log.warning('Quantization: nncf installed please restart')
+    install('jstyleson', quiet=True)
+    install('texttable', quiet=True)
     try:
         import nncf
         intel_nncf = nncf
@@ -350,6 +352,7 @@ def nncf_compress_model(model, op=None, sd_model=None, send_to_device=True, do_g
     num_bits = 8 if shared.opts.nncf_compress_weights_mode in {"INT8", "INT8_SYM", "INT8_ASYM"} else 4
     is_asym_mode = shared.opts.nncf_compress_weights_mode in {"INT8", "INT4", "INT8_ASYM", "INT4_ASYM"}
     model = apply_nncf_to_module(model, num_bits, is_asym_mode, quant_conv=shared.opts.nncf_quantize_conv_layers)
+    model.quantization_method = 'NNCF'
     if send_to_device:
         nncf_send_to_device(model, devices.device)
 

modules/modelloader.py

@@ -26,7 +26,7 @@ def hf_login(token=None):
     global loggedin # pylint: disable=global-statement
     token = token or shared.opts.huggingface_token
     install('hf_xet', quiet=True)
-    if token is None or len(token) <= 2:
+    if token is None or len(token) <= 4:
         log.debug('HF login: no token provided')
         return False
     if os.environ.get('HUGGING_FACE_HUB_TOKEN', None) is not None:
@@ -41,8 +41,9 @@ def hf_login(token=None):
             hf.logout()
             hf.login(token=token, add_to_git_credential=False, write_permission=False)
         text = stdout.getvalue() or ''
+        obfuscated_token = 'hf_...' + token[-4:]
         line = [l for l in text.split('\n') if 'Token' in l]
-        log.info(f'HF login: token="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}')
+        log.info(f'HF login: token="{obfuscated_token}" fn="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}')
         loggedin = token
     return True
 

modules/processing_args.py

@@ -35,7 +35,10 @@ def task_specific_kwargs(p, model):
             }
     elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
         if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'):
-            model.register_to_config(requires_aesthetics_score = False)
+            if model.__class__.__name__ in sd_models.i2i_pipes:
+                pass
+            else:
+                model.register_to_config(requires_aesthetics_score = False)
         if 'hires' not in p.ops:
             p.ops.append('img2img')
         if p.vae_type == 'Remote':
@@ -71,7 +74,10 @@ def task_specific_kwargs(p, model):
         }
     elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
         if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'):
-            model.register_to_config(requires_aesthetics_score = False)
+            if model.__class__.__name__ in [sd_models.i2i_pipes]:
+                pass
+            else:
+                model.register_to_config(requires_aesthetics_score = False)
         if p.detailer_enabled:
             p.ops.append('detailer')
         else:
@@ -144,7 +150,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
         'StableDiffusion' in model.__class__.__name__ or
         'StableCascade' in model.__class__.__name__ or
         'Flux' in model.__class__.__name__ or
-        'HiDreamImage' in model.__class__.__name__
+        'HiDreamImagePipeline' in model.__class__.__name__ # hidream-e1 has different embeds
     ):
         try:
             prompt_parser_diffusers.embedder = prompt_parser_diffusers.PromptEmbedder(prompts, negative_prompts, steps, clip_skip, p)
@@ -161,7 +167,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     if 'prompt' in possible:
         if 'OmniGen' in model.__class__.__name__:
             prompts = [p.replace('|image|', '<|image_1|>') for p in prompts]
-        if 'HiDreamImage' in model.__class__.__name__:
+        if 'HiDreamImage' in model.__class__.__name__  and prompt_parser_diffusers.embedder is not None:
             args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
             prompt_embeds = prompt_parser_diffusers.embedder('prompt_embeds')
             args['prompt_embeds_t5'] = prompt_embeds[0]
@@ -180,7 +186,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
         else:
             args['prompt'] = prompts
     if 'negative_prompt' in possible:
-        if 'HiDreamImage' in model.__class__.__name__:
+        if 'HiDreamImage' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
             args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
             negative_prompt_embeds = prompt_parser_diffusers.embedder('negative_prompt_embeds')
             args['negative_prompt_embeds_t5'] = negative_prompt_embeds[0]

modules/processing_diffusers.py

@@ -54,7 +54,7 @@ def restore_state(p: processing.StableDiffusionProcessing):
 
 def process_base(p: processing.StableDiffusionProcessing):
     txt2img = is_txt2img()
-    use_refiner_start = txt2img and is_refiner_enabled(p) and not p.is_hr_pass and p.refiner_start > 0 and p.refiner_start < 1
+    use_refiner_start = is_refiner_enabled(p) and (not p.is_hr_pass)
     use_denoise_start = not txt2img and p.refiner_start > 0 and p.refiner_start < 1
 
     shared.sd_model = update_pipeline(shared.sd_model, p)

modules/processing_helpers.py

@@ -23,7 +23,7 @@ def is_txt2img():
 
 
 def is_refiner_enabled(p):
-    return p.enable_hr and p.refiner_steps > 0 and p.refiner_start > 0 and p.refiner_start < 1 and shared.sd_refiner is not None
+    return p.enable_hr and (p.refiner_steps > 0) and (p.refiner_start > 0) and (p.refiner_start < 1) and (shared.sd_refiner is not None)
 
 
 def setup_color_correction(image):
@@ -454,12 +454,15 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start):
     if len(getattr(p, 'timesteps', [])) > 0:
         return None
     if not is_txt2img():
-        if use_denoise_start and shared.sd_model_type == 'sdxl':
+        cls = shared.sd_model.__class__.__name__
+        if cls in sd_models.i2i_pipes:
+            steps = p.steps
+        elif 'Flex' in cls:
+            steps = p.steps
+        elif 'HiDreamImageEditingPipeline' in cls:
+            steps = p.steps
+        elif use_denoise_start and (shared.sd_model_type == 'sdxl'):
             steps = p.steps // (1 - p.refiner_start)
-        elif 'Flex' in shared.sd_model.__class__.__name__:
-            steps = p.steps
-        elif shared.sd_model_type == 'omnigen':
-            steps = p.steps
         elif p.denoising_strength > 0:
             steps = (p.steps // p.denoising_strength) + 1
         else:
@@ -536,19 +539,17 @@ def set_latents(p):
     return latents
 
 
-last_circular = False
-def apply_circular(enable, model):
-    global last_circular # pylint: disable=global-statement
+def apply_circular(enable: bool, model):
     if not hasattr(model, 'unet') or not hasattr(model, 'vae'):
         return
-    if last_circular == enable:
+    if getattr(model, 'texture_tiling', False) == enable:
         return
     try:
         for layer in [layer for layer in model.unet.modules() if type(layer) is torch.nn.Conv2d]:
             layer.padding_mode = 'circular' if enable else 'zeros'
         for layer in [layer for layer in model.vae.modules() if type(layer) is torch.nn.Conv2d]:
             layer.padding_mode = 'circular' if enable else 'zeros'
-        last_circular = enable
+        model.texture_tiling = enable
     except Exception as e:
         debug(f"Diffusers tiling failed: {e}")
 

modules/sd_checkpoint.py

@@ -158,6 +158,7 @@ def list_models():
     shared.log.info(f'Available Models: safetensors="{shared.opts.ckpt_dir}":{len(safetensors_list)} diffusers="{shared.opts.diffusers_dir}":{len(diffusers_list)} items={len(checkpoints_list)} time={time.time()-t0:.2f}')
     checkpoints_list = dict(sorted(checkpoints_list.items(), key=lambda cp: cp[1].filename))
 
+
 def update_model_hashes():
     txt = []
     lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None]

modules/sd_detect.py

@@ -104,12 +104,14 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 index = shared.readfile(index, silent=True)
                 cls = index.get('_class_name', None)
                 if cls is not None:
-                    pipeline = getattr(diffusers, cls)
-                if 'Flux' in pipeline.__name__ and guess != 'FLEX':
+                    pipeline = getattr(diffusers, cls, None)
+                    if pipeline is None:
+                        pipeline = cls
+                if callable(pipeline) and 'Flux' in pipeline.__name__ and guess != 'FLEX':
                     guess = 'FLUX'
-                if 'StableDiffusion3' in pipeline.__name__:
+                if callable(pipeline) and 'StableDiffusion3' in pipeline.__name__:
                     guess = 'Stable Diffusion 3'
-                if 'Lumina2' in pipeline.__name__:
+                if callable(pipeline) and 'Lumina2' in pipeline.__name__:
                     guess = 'Lumina 2'
             # switch for specific variant
             if guess == 'Stable Diffusion' and 'inpaint' in f.lower():

modules/sd_models.py

@@ -31,19 +31,27 @@ debug_process = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not Non
 diffusers_version = int(diffusers.__version__.split('.')[1])
 checkpoint_tiles = checkpoint_titles # legacy compatibility
 pipe_switch_task_exclude = [
-    'StableDiffusionReferencePipeline',
-    'StableDiffusionAdapterPipeline',
-    'AnimateDiffPipeline',
-    'AnimateDiffSDXLPipeline',
-    'OmniGenPipeline',
-    'StableDiffusion3ControlNetPipeline',
-    'InstantIRPipeline',
-    'FluxFillPipeline',
+    'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline',
     'FluxControlPipeline',
-    'PixelSmithXLPipeline',
-    'PhotoMakerStableDiffusionXLPipeline',
-    'StableDiffusionXLInstantIDPipeline',
+    'FluxFillPipeline',
+    'InstantIRPipeline',
     'LTXConditionPipeline',
+    'OmniGenPipeline',
+    'PhotoMakerStableDiffusionXLPipeline',
+    'PixelSmithXLPipeline',
+    'StableDiffusion3ControlNetPipeline',
+    'StableDiffusionAdapterPipeline',
+    'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline',
+    'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
+    'StableDiffusionReferencePipeline',
+    'StableDiffusionXLInstantIDPipeline',
+]
+i2i_pipes = [
+    'LEditsPPPipelineStableDiffusion',
+    'LEditsPPPipelineStableDiffusionXL',
+    'OmniGenPipeline',
+    'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline',
+    'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
 ]
 
 
@@ -276,7 +284,7 @@ def load_diffuser_initial(diffusers_load_config, op='model'):
 
 def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='model'):
     sd_model = None
-    unload_model_weights()
+    unload_model_weights(op=op)
     shared.sd_model = None
     try:
         if model_type in ['Stable Cascade']: # forced pipeline
@@ -655,7 +663,7 @@ class DiffusersTaskType(Enum):
 
 def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
     cls = pipe.__class__.__name__
-    if cls in ["LEditsPPPipelineStableDiffusion", "LEditsPPPipelineStableDiffusionXL", "OmniGenPipeline"]: # special case
+    if cls in i2i_pipes: # special case
         return DiffusersTaskType.IMAGE_2_IMAGE
     elif 'ImageToVideo' in cls or cls in ['LTXConditionPipeline', 'StableVideoDiffusionPipeline']: # i2v pipelines
         return DiffusersTaskType.IMAGE_2_IMAGE
@@ -782,6 +790,9 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE:
         clean_diffuser_pipe(pipe)
 
+    if hasattr(pipe, 'no_task_switch'):
+        del pipe.no_task_switch
+        return pipe
     if get_diffusers_task(pipe) == new_pipe_type:
         return pipe
 

modules/sd_samplers_diffusers.py

@@ -232,6 +232,8 @@ class DiffusionSampler:
                 self.config['beta_schedule'] = 'scaled_linear'
             elif shared.opts.schedulers_beta_schedule == 'cosine':
                 self.config['beta_schedule'] = 'squaredcos_cap_v2'
+            elif shared.opts.schedulers_beta_schedule == 'sigmoid':
+                self.config['beta_schedule'] = 'sigmoid'
 
         timesteps = re.split(',| ', shared.opts.schedulers_timesteps)
         timesteps = [int(x) for x in timesteps if x.isdigit()]

modules/sd_vae_remote.py

@@ -28,13 +28,6 @@ dtypes = {
 }
 
 
-def h1_pack_latents(latents, _batch_size, _num_channels_latents, _height, _width): # TODO hidream: pack latents for remote vae
-    # latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
-    # latents = latents.permute(0, 2, 4, 1, 3, 5)
-    # latents = latents.reshape(batch_size, (height // 2) * (width // 2) // (num_channels_latents * 4), num_channels_latents * 4)
-    return latents
-
-
 def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_type: str = None) -> Image.Image:
     from modules import devices, shared, errors, modelloader
     tensors = []
@@ -57,9 +50,6 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
             latent = latent_copy[i]
             if model_type != 'f1':
                 latent = latent.unsqueeze(0)
-            # if model_type == 'h1':
-            #     num_channels_latents = shared.sd_model.transformer.config.in_channels
-            #     latent = h1_pack_latents(latent, 1, num_channels_latents, height, width) # pylint: disable=protected-access
             params = {
                 "input_tensor_type": "binary",
                 "shape": list(latent.shape),

modules/shared.py

@@ -814,7 +814,7 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"),
     "schedulers_sigma": OptionInfo("default", "Sigma algorithm", gr.Radio, {"choices": ['default', 'karras', 'exponential', 'polyexponential'], "visible": False}),
 
     # managed from ui.py for backend diffusers
-    "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2'], "visible": False}),
+    "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2', 'sigmoid'], "visible": False}),
     "schedulers_use_thresholding": OptionInfo(False, "Use dynamic thresholding", gr.Checkbox, {"visible": False}),
     "schedulers_timestep_spacing": OptionInfo("default", "Timestep spacing", gr.Dropdown, {"choices": ['default', 'linspace', 'leading', 'trailing'], "visible": False}),
     'schedulers_timesteps': OptionInfo('', "Timesteps", gr.Textbox, {"visible": False}),
@@ -973,7 +973,6 @@ options_templates.update(options_section(('extra_networks', "Networks"), {
 
     "extra_networks_styles_sep": OptionInfo("<h2>Styles</h2>", "", gr.HTML),
     "extra_networks_styles": OptionInfo(True, "Show reference styles"),
-    "extra_networks_save_unparsed": OptionInfo(True, "Save unparsed prompt"),
     "extra_networks_apply_unparsed": OptionInfo(True, "Restore unparsed prompt"),
 
     "extra_networks_embed_sep": OptionInfo("<h2>Embeddings</h2>", "", gr.HTML),

modules/shared_items.py

@@ -17,6 +17,7 @@ pipelines = {
     'Stable Diffusion XL': getattr(diffusers, 'StableDiffusionXLPipeline', None),
     'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None),
     'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None),
+    'Stable Diffusion XL Refiner': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None),
     'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None),
     'Stable Diffusion 3.x': getattr(diffusers, 'StableDiffusion3Pipeline', None),
     'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None),

modules/styles.py

@@ -152,11 +152,6 @@ def apply_styles_to_extra(p, style: Style):
     reference_style = get_reference_style()
     extra = infotext.parse(reference_style) if shared.opts.extra_network_reference_values else {}
 
-    if not hasattr(p, 'original_prompt'):
-        p.original_prompt = p.prompt
-    if not hasattr(p, 'original_negative'):
-        p.original_negative = p.negative_prompt
-
     style_extra = apply_wildcards_to_prompt(style.extra, [style.wildcards], silent=True)
     style_extra = ' ' + style_extra.lower()
     extra.update(infotext.parse(style_extra))
@@ -340,6 +335,11 @@ class StyleDatabase:
         return prompt
 
     def apply_styles_to_extra(self, p):
+        if len(getattr(p, 'original_prompt', '')) == 0:
+            p.original_prompt = p.prompt
+        if len(getattr(p, 'original_negative', '')) == 0:
+            p.original_negative = p.negative_prompt
+
         if p.styles is None:
             return
         if p.styles is None or not isinstance(p.styles, list):

modules/ui_control.py

@@ -359,7 +359,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             controlnetxs_ui_units.append(unit_ui)
                             units.append(unit.Unit(
                                 unit_type = 'xs',
-                                index = 1,
+                                index = i,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,

modules/ui_extra_networks.py

@@ -783,8 +783,10 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
 
     def show_details(text, img, desc, info, meta, description, prompt, negative, parameters, wildcards, params, _dummy1=None, _dummy2=None):
         page, item = get_item(state, params)
-        valid = item is not None and hasattr(item, 'name') and hasattr(item, 'filename')
-        if valid:
+        is_style = (page is not None) and (page.title == 'Style')
+        is_valid = (item is not None) and hasattr(item, 'name') and hasattr(item, 'filename')
+
+        if is_valid:
             stat = os.stat(item.filename) if os.path.exists(item.filename) else None
             desc = item.description
             fullinfo = shared.readfile(os.path.splitext(item.filename)[0] + '.json', silent=True)
@@ -795,16 +797,14 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
                 item.filename = None
                 shared.log.warning('Network: show details not supported for compound item')
                 info = None
-            """
-            if prompt is not None:
+            if prompt is not None and len(prompt) > 0:
                 item.prompt = prompt
-            if negative is not None:
+            if negative is not None and len(negative) > 0:
                 item.negative = negative
-            if description is not None:
+            if description is not None and len(description) > 0:
                 item.description = description
-            if wildcards is not None:
+            if wildcards is not None and len(wildcards) > 0:
                 item.wildcards = wildcards
-            """
             meta = page.metadata.get(item.name, {}) or {}
             if type(meta) is str:
                 try:
@@ -887,7 +887,6 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
                 </tbody></table>
                 {note}
             '''
-        is_style = (page is not None) and (page.title == 'Style')
         return [
             text, # gr.html
             img, # gr.image
@@ -899,7 +898,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
             gr.update(value=negative, visible=is_style), # gr.textbox
             gr.update(value=parameters, visible=is_style), # gr.textbox
             gr.update(value=wildcards, visible=is_style), # gr.textbox
-            gr.update(visible=valid), # details ui visible
+            gr.update(visible=is_valid), # details ui visible
             gr.update(visible=not is_style), # details ui tabs visible
             gr.update(visible=is_style), # details ui text visible
         ]
@@ -940,10 +939,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         return ui_refresh_click(title)
 
     def ui_save_click():
-        if shared.opts.extra_networks_save_unparsed:
-            from modules.processing_info import get_last_args
-            params, text = get_last_args()
-        else:
+        from modules.processing_info import get_last_args
+        params, text = get_last_args()
+        if (not params) or (not text) or (len(text) == 0):
             filename = os.path.join(paths.data_path, "params.txt")
             if os.path.exists(filename):
                 with open(filename, "r", encoding="utf8") as file:
@@ -957,13 +955,12 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         return res
 
     def ui_quicksave_click(name):
-        if shared.opts.extra_networks_save_unparsed:
-            from modules.processing_info import get_last_args
-            params, text = get_last_args()
-        else:
-            if name is None or len(name) < 1:
-                shared.log.warning("Network quick save style: no name provided")
-                return
+        if name is None or len(name) < 1:
+            shared.log.warning("Network quick save style: no name provided")
+            return
+        from modules.processing_info import get_last_args
+        params, text = get_last_args()
+        if (not params) or (not text) or (len(text) == 0):
             fn = os.path.join(paths.data_path, "params.txt")
             if os.path.exists(fn):
                 with open(fn, "r", encoding="utf8") as file:
@@ -983,9 +980,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         }
         shared.writefile(item, fn, silent=True)
         if len(prompt) > 0:
-            shared.log.debug(f'Network quick save style: item="{name}" filename="{fn}" unparsed={shared.opts.extra_networks_unparsed}')
+            shared.log.debug(f'Networks type=style quicksave style: item="{name}" filename="{fn}" prompt="{prompt}"')
         else:
-            shared.log.warning(f'Network quick save model: item="{name}" filename="{fn}" prompt is empty')
+            shared.log.warning(f'Networks type=style quicksave model: item="{name}" filename="{fn}" prompt is empty')
 
     def ui_sort_cards(sort_order):
         if shared.opts.extra_networks_sort != sort_order:

modules/ui_sections.py

@@ -307,7 +307,7 @@ def create_sampler_options(tabname):
             sampler_sigma = gr.Dropdown(label='Sigma method', elem_id=f"{tabname}_sampler_sigma", choices=['default', 'karras', 'betas', 'exponential', 'lambdas', 'flowmatch'], value=shared.opts.schedulers_sigma, type='value')
             sampler_spacing = gr.Dropdown(label='Timestep spacing', elem_id=f"{tabname}_sampler_spacing", choices=['default', 'linspace', 'leading', 'trailing'], value=shared.opts.schedulers_timestep_spacing, type='value')
         with gr.Row(elem_classes=['flex-break']):
-            sampler_beta = gr.Dropdown(label='Beta schedule', elem_id=f"{tabname}_sampler_beta", choices=['default', 'linear', 'scaled', 'cosine'], value=shared.opts.schedulers_beta_schedule, type='value')
+            sampler_beta = gr.Dropdown(label='Beta schedule', elem_id=f"{tabname}_sampler_beta", choices=['default', 'linear', 'scaled', 'cosine', 'sigmoid'], value=shared.opts.schedulers_beta_schedule, type='value')
             sampler_prediction = gr.Dropdown(label='Prediction method', elem_id=f"{tabname}_sampler_prediction", choices=['default', 'epsilon', 'sample', 'v_prediction', 'flow_prediction'], value=shared.opts.schedulers_prediction_type, type='value')
         with gr.Row(elem_classes=['flex-break']):
             sampler_presets = gr.Dropdown(label='Timesteps presets', elem_id=f"{tabname}_sampler_presets", choices=['None', 'AYS SD15', 'AYS SDXL'], value='None', type='value')

modules/ui_video_vlm.py

@@ -0,0 +1,66 @@
+import gradio as gr
+from PIL import Image
+from modules import shared
+
+
+models = [
+    'Google Gemma 3 4B',
+]
+system_prompts = {
+    'prefix': 'You are a video director.',
+    'desc': 'You will provide a prompt that can be used to create short, creative and dynamic video. Describe the subject of the video, add movement that fits the overall tone of the video, add camera movement.',
+    'suffix': 'NSFW content and nudity are allowed and must be fully described if present. Movement should be dynamic and creative. Do not specify duration and assume video is short. Avoid slow-motion and prefer faster movements. Output should be a single short paragraph without explanations',
+    'example': 'Example: "Short video of beautiful blonde woman in her 20ies wearing a long flowing red dress. She is briskly walking on the beach during sunset and performing a pirouette ending with her hand pointing at the camera as she smiles. Camera is moving around her and zooming to her face. Sun is setting in the background causing changes in colors and shadows to move dynamically."',
+
+    't2v-prompt': 'You are a given short prompt with basic instructions.',
+    't2v-noprompt': '',
+    'i2v-prompt': 'You are given an image as a starting point and a short prompt with basic instructions.',
+    'i2v-noprompt': 'You are given an image as a starting point.',
+}
+
+
+def enhance_prompt(enable:bool, model:str=None, image=None, prompt:str='', system_prompt:str=''):
+    from modules.interrogate import vqa
+    if not enable:
+        return prompt
+    if model is None or len(model) < 4:
+        model = models[0]
+    if image is not None and not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    if prompt is None or len(prompt) < 4:
+        prompt = '  '
+    if system_prompt is None or len(system_prompt) < 4:
+        if image is not None:
+            if prompt is not None and len(prompt) > 4:
+                core_prompt = system_prompts['i2v-prompt']
+            else:
+                core_prompt = system_prompts['i2v-noprompt']
+        else:
+            if prompt is not None and len(prompt) > 4:
+                core_prompt = system_prompts['t2v-prompt']
+            else:
+                core_prompt = system_prompts['t2v-noprompt']
+        system_prompt = f"{system_prompts['prefix']} {core_prompt} {system_prompts['desc']} {system_prompts['suffix']} {system_prompts['example']}"
+    shared.log.debug(f'Video prompt enhance: model="{model}" image={image} prompt="{prompt}"')
+    # shared.log.trace(f'Video prompt enhance: system="{system_prompt}"')
+    answer = vqa.interrogate(question='', prompt=prompt, system_prompt=system_prompt, image=image, model_name=model, quiet=False)
+    shared.log.debug(f'Video prompt enhance: answer="{answer}"')
+    return answer
+
+
+def create_ui(prompt_element:gr.Textbox, image_element:gr.Image):
+    with gr.Accordion('Prompt enhance', open=False):
+        with gr.Row():
+            enable = gr.Checkbox(label='Enable', value=False)
+            btn_enhance = gr.Button(value='Enhance now', elem_id='btn_enhance')
+        with gr.Row():
+            model = gr.Dropdown(label='Model', choices=models, value=models[0])
+        with gr.Row():
+            system_prompt = gr.Textbox(label='System prompt', placeholder='override system prompt with user-provided prompt', lines=3)
+        btn_enhance.click(
+            fn=enhance_prompt,
+            inputs=[enable, model, image_element, prompt_element, system_prompt],
+            outputs=prompt_element,
+            show_progress=True,
+        )
+    return enable, model, system_prompt

modules/zluda_hijacks.py

@@ -33,16 +33,8 @@ MEM_BUS_WIDTH = {
 }
 
 
-_topk = torch.topk
-def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-builtin
-    device = input.device
-    values, indices = _topk(input.cpu(), *args, **kwargs)
-    return torch.return_types.topk((values.to(device), indices.to(device),))
-
-
 class DeviceProperties:
     PROPERTIES_OVERRIDE = {
-        "regs_per_multiprocessor": 65535,
          # sometimes gcnArchName contains device name ("AMD Radeon RX ..."), not architecture name ("gfx...")
         "gcnArchName": "UNKNOWN ARCHITECTURE",
     }
@@ -68,7 +60,6 @@ def torch__C__cuda_getCurrentRawStream(device):
 
 
 def do_hijack():
-    torch.topk = topk
     if zluda.default_agent is not None:
         DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name
     torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access
@@ -104,10 +95,13 @@ def do_hijack():
                         query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
                         key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
                         value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
-                    out_padded = interface_fa.fwd(
-                        query.transpose(1, 2),
+                    query = query.transpose(1, 2)
+                    out_padded = torch.zeros_like(query)
+                    interface_fa.fwd(
+                        query,
                         key.transpose(1, 2),
                         value.transpose(1, 2),
+                        out_padded,
                         dropout_p,
                         scale,
                         is_causal,

modules/zluda_installer.py

@@ -69,8 +69,8 @@ def set_default_agent(agent: rocm.Agent):
     default_agent = agent
 
 
-def is_reinstall_needed() -> bool: # ZLUDA<3.8.7
-    return not os.path.exists(os.path.join(path, 'cufftw.dll'))
+def is_reinstall_needed() -> bool: # ZLUDA<3.9.4
+    return os.path.exists(os.path.join(path, 'cudart.dll'))
 
 
 def install():
@@ -78,7 +78,7 @@ def install():
         return
 
     platform = "windows"
-    commit = os.environ.get("ZLUDA_HASH", "dba64c0966df2c71e82255e942c96e2e1cea3a2d")
+    commit = os.environ.get("ZLUDA_HASH", "8d2128caf460b853b165cab0b4d8826b6b734ae7")
     if os.environ.get("ZLUDA_NIGHTLY", "0") == "1":
         log.warning("Environment variable 'ZLUDA_NIGHTLY' will be removed. Please use command-line argument '--use-nightly' instead.")
         args.use_nightly = True

requirements.txt

@@ -31,7 +31,7 @@ invisible-watermark
 pi-heif
 
 # versioned
-rich==13.9.4
+rich==14.0.0
 safetensors==0.5.3
 tensordict==0.1.2
 peft==0.15.2
@@ -46,9 +46,9 @@ opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
 huggingface_hub==0.30.2
-numexpr==2.8.8
+numexpr==2.10.2
 numpy==1.26.4
-numba==0.59.1
+numba==0.61.2
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.21.1

scripts/prompt_enhance.py

@@ -16,6 +16,12 @@ class Options:
     models = {
         'google/gemma-3-1b-it': {},
         'google/gemma-3-4b-it': {},
+        'Qwen/Qwen3-0.6B-FP8': {},
+        'Qwen/Qwen3-1.7B-FP8': {},
+        'Qwen/Qwen3-4B-FP8': {},
+        'Qwen/Qwen3-0.6B': {},
+        'Qwen/Qwen3-1.7B': {},
+        'Qwen/Qwen3-4B': {},
         'Qwen/Qwen2.5-0.5B-Instruct': {},
         'Qwen/Qwen2.5-1.5B-Instruct': {},
         'Qwen/Qwen2.5-3B-Instruct': {},
@@ -50,6 +56,7 @@ class Options:
     do_sample: bool = True
     temperature: float = 0.15
     repetition_penalty: float = 1.2
+    thinking_mode: bool = False
 
 
 class Script(scripts.Script):
@@ -195,7 +202,7 @@ class Script(scripts.Script):
         filtered = re.sub(pattern, '', prompt)
         return filtered, matches
 
-    def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None):
+    def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None, thinking:bool=False):
         model = model or self.options.default
         prompt = prompt or self.prompt.value
         prefix = prefix or ''
@@ -204,6 +211,7 @@ class Script(scripts.Script):
         tokens = tokens or self.options.max_tokens
         penalty = penalty or self.options.repetition_penalty
         temperature = temperature or self.options.temperature
+        thinking = thinking or self.options.thinking_mode
         sample = sample if sample is not None else self.options.do_sample
         while self.busy:
             time.sleep(0.1)
@@ -223,6 +231,7 @@ class Script(scripts.Script):
             inputs = self.tokenizer.apply_chat_template(
                 chat_template,
                 add_generation_prompt=True,
+                enable_thinking=thinking,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
@@ -270,7 +279,7 @@ class Script(scripts.Script):
             response = self.post(response, prefix, suffix, networks)
         shared.log.info(f'Prompt enhance: model="{model}" time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt={len(prompt)} response={len(response)}')
         if debug_enabled:
-            shared.log.trace(f'Prompt enhance: sample={sample} tokens={tokens} temperature={temperature} penalty={penalty}')
+            shared.log.trace(f'Prompt enhance: sample={sample} tokens={tokens} temperature={temperature} penalty={penalty} thinking={thinking}')
             shared.log.trace(f'Prompt enhance: prompt="{prompt}"')
             shared.log.trace(f'Prompt enhance: response="{response}"')
         self.busy = False
@@ -279,7 +288,7 @@ class Script(scripts.Script):
             return prompt
         return response
 
-    def apply(self, prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty):
+    def apply(self, prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode):
         response = self.enhance(
             prompt=prompt,
             prefix=prompt_prefix,
@@ -290,6 +299,7 @@ class Script(scripts.Script):
             tokens=max_tokens,
             temperature=temperature,
             penalty=repetition_penalty,
+            thinking=thinking_mode,
         )
         if apply_prompt:
             return [response, response]
@@ -339,6 +349,8 @@ class Script(scripts.Script):
                     with gr.Row():
                         temperature = gr.Slider(label='Temperature', value=self.options.temperature, minimum=0.0, maximum=1.0, step=0.01, interactive=True)
                         repetition_penalty = gr.Slider(label='Repetition penalty', value=self.options.repetition_penalty, minimum=0.0, maximum=2.0, step=0.01, interactive=True)
+                    with gr.Row():
+                        thinking_mode = gr.Checkbox(label='Thinking mode', value=False, interactive=True)
                     gr.HTML('<br>')
                 with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'):
                     with gr.Row():
@@ -355,15 +367,15 @@ class Script(scripts.Script):
                         clear_btn.click(fn=lambda: '', inputs=[], outputs=[prompt_output])
                         copy_btn = gr.Button(value='Set prompt', elem_id='prompt_enhance_copy', variant='secondary')
                         copy_btn.click(fn=lambda x: x, inputs=[prompt_output], outputs=[self.prompt])
-            apply_btn.click(fn=self.apply, inputs=[self.prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty], outputs=[prompt_output, self.prompt])
-        return [apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty]
+            apply_btn.click(fn=self.apply, inputs=[self.prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode], outputs=[prompt_output, self.prompt])
+        return [apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode]
 
     def after_component(self, component, **kwargs): # searching for actual ui prompt components
         if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt', 'video_prompt']:
             self.prompt = component
 
     def before_process(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=unused-argument
-        apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty = args
+        apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode = args
         if not apply_auto and not p.enhance_prompt:
             return
         if shared.state.skipped or shared.state.interrupted:
@@ -383,6 +395,7 @@ class Script(scripts.Script):
             tokens=max_tokens,
             temperature=temperature,
             penalty=repetition_penalty,
+            thinking=thinking_mode,
         )
         p.extra_generation_params['LLM'] = llm_model
         shared.state.end()

scripts/pulid_ext.py

@@ -137,7 +137,7 @@ class Script(scripts.Script):
             shared.log.error('PuLID: no images')
             return None
 
-        supported_model_list = ['sdxl']
+        supported_model_list = ['sdxl', 'f1']
         if shared.sd_model_type not in supported_model_list:
             shared.log.error(f'PuLID: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}')
             return None
@@ -168,15 +168,12 @@ class Script(scripts.Script):
             p.batch_size = 1
 
         sdp = shared.opts.cross_attention_optimization == "Scaled-Dot-Product"
-        sampler_fn = getattr(self.pulid.sampling, f'sample_{sampler}', None)
         strength = getattr(p, 'pulid_strength', strength)
         zero = getattr(p, 'pulid_zero', zero)
         ortho = getattr(p, 'pulid_ortho', ortho)
         sampler = getattr(p, 'pulid_sampler', sampler)
         restore = getattr(p, 'pulid_restore', restore)
         p.pulid_restore = restore
-        if sampler_fn is None:
-            sampler_fn = self.pulid.sampling.sample_dpmpp_2m_sde
 
         if shared.sd_model_type == 'sdxl' and not hasattr(shared.sd_model, 'pipe'):
             try:
@@ -203,9 +200,42 @@ class Script(scripts.Script):
                 shared.log.error(f'PuLID: failed to create pipeline: {e}')
                 errors.display(e, 'PuLID')
                 return None
+        elif shared.sd_model_type == 'f1':
+            # TODO nunchaku: pulid-f1
+            shared.log.error('PuLID: f1 not supported')
+            return None
 
+        if shared.sd_model_type == 'sdxl':
+            processed = self.run_sdxl(p, images, strength, zero, sampler, ortho, restore, offload, version)
+        elif shared.sd_model_type == 'f1':
+            processed = None
+        else:
+            shared.log.error(f'PuLID: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}')
+            processed = None
+        return processed
+
+    def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args): # pylint: disable=unused-argument
+        _strength, _zero, _sampler, _ortho, _gallery, restore, _offload, _version = args
+        if shared.sd_model_type == "sdxl" and hasattr(shared.sd_model, 'pipe'):
+            restore = getattr(p, 'pulid_restore', restore)
+            if restore:
+                if hasattr(shared.sd_model, 'app'):
+                    shared.sd_model.app = None
+                    shared.sd_model.ip_adapter = None
+                    shared.sd_model.face_helper = None
+                    shared.sd_model.clip_vision_model = None
+                    shared.sd_model.handler_ante = None
+                shared.sd_model = shared.sd_model.pipe
+                devices.torch_gc(force=True)
+            shared.log.debug(f'PuLID complete: class={shared.sd_model.__class__.__name__} preprocess={self.preprocess:.2f} pipe={"restore" if restore else "cache"}')
+        return processed
+
+    def run_sdxl(self, p: processing.StableDiffusionProcessing, images: list, strength: float, zero: int, sampler: str, ortho: str, restore: bool, offload: bool, version: str):
+        sampler_fn = getattr(self.pulid.sampling, f'sample_{sampler}', None)
+        if sampler_fn is None:
+            sampler_fn = self.pulid.sampling.sample_dpmpp_2m_sde
         shared.sd_model.sampler = sampler_fn
-        shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} version="{version}" sdp={sdp} strength={strength} zero={zero} ortho={ortho} sampler={sampler_fn} images={[i.shape for i in images]} offload={offload} restore={restore}')
+        shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} version="{version}" strength={strength} zero={zero} ortho={ortho} sampler={sampler_fn} images={[i.shape for i in images]} offload={offload} restore={restore}')
         self.pulid.attention.NUM_ZERO = zero
         self.pulid.attention.ORTHO = ortho == 'v1'
         self.pulid.attention.ORTHO_v2 = ortho == 'v2'
@@ -256,19 +286,3 @@ class Script(scripts.Script):
         # interim = [Image.fromarray(img) for img in shared.sd_model.debug_img_list]
         # shared.log.debug(f'PuLID: time={t1-t0:.2f}')
         return processed
-
-    def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args): # pylint: disable=unused-argument
-        _strength, _zero, _sampler, _ortho, _gallery, restore, _offload, _version = args
-        if hasattr(shared.sd_model, 'pipe') and shared.sd_model_type == "sdxl":
-            restore = getattr(p, 'pulid_restore', restore)
-            if restore:
-                if hasattr(shared.sd_model, 'app'):
-                    shared.sd_model.app = None
-                    shared.sd_model.ip_adapter = None
-                    shared.sd_model.face_helper = None
-                    shared.sd_model.clip_vision_model = None
-                    shared.sd_model.handler_ante = None
-                shared.sd_model = shared.sd_model.pipe
-                devices.torch_gc(force=True)
-            shared.log.debug(f'PuLID complete: class={shared.sd_model.__class__.__name__} preprocess={self.preprocess:.2f} pipe={"restore" if restore else "cache"}')
-        return processed