Upload src/pipeline.py with huggingface_hub

src/pipeline.py

@@ -103,6 +103,66 @@ def _build_fused_qkv_weight(q_w, k_w, v_w) -> torch.Tensor:
 # TP-aware attention processors
 # ---------------------------------------------------------------------------
 
+class Flux2AttnProcessorFlashAttn(Flux2AttnProcessor):
+    """
+    TP-aware double-stream attention processor using a NKI flash attention kernel.
+    Replaces dispatch_attention_fn with flux2_flash_attn (BSHD in, BSHD out).
+    """
+    def __init__(self, tp_size: int):
+        super().__init__()
+        self.tp_size = tp_size
+        self._flash_attn = None
+
+    def _init_nki(self):
+        if self._flash_attn is None:
+            import sys, os
+            _dir = os.path.dirname(os.path.abspath(__file__))
+            if _dir not in sys.path:
+                sys.path.insert(0, _dir)
+            from nki_flash_attn import flux2_flash_attn
+            self._flash_attn = flux2_flash_attn
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None,
+                 attention_mask=None, image_rotary_emb=None, **kwargs):
+        self._init_nki()
+        local_heads = attn.heads // self.tp_size
+        head_dim    = attn.head_dim
+
+        query, key, value, encoder_query, encoder_key, encoder_value = \
+            _get_qkv_projections(attn, hidden_states, encoder_hidden_states)
+        query = query.unflatten(-1, (local_heads, head_dim))
+        key   = key.unflatten(-1,   (local_heads, head_dim))
+        value = value.unflatten(-1, (local_heads, head_dim))
+        query = attn.norm_q(query)
+        key   = attn.norm_k(key)
+        if attn.added_kv_proj_dim is not None:
+            encoder_query = encoder_query.unflatten(-1, (local_heads, head_dim))
+            encoder_key   = encoder_key.unflatten(-1,   (local_heads, head_dim))
+            encoder_value = encoder_value.unflatten(-1, (local_heads, head_dim))
+            encoder_query = attn.norm_added_q(encoder_query)
+            encoder_key   = attn.norm_added_k(encoder_key)
+            query = torch.cat([encoder_query, query], dim=1)
+            key   = torch.cat([encoder_key,   key],   dim=1)
+            value = torch.cat([encoder_value, value], dim=1)
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key   = apply_rotary_emb(key,   image_rotary_emb, sequence_dim=1)
+
+        hidden_states = self._flash_attn(query, key, value)  # (B, S, H_local, D) BSHD
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+                [encoder_hidden_states.shape[1],
+                 hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+        if encoder_hidden_states is not None:
+            return hidden_states, encoder_hidden_states
+        return hidden_states
+
+
 class Flux2AttnProcessorTP(Flux2AttnProcessor):
     def __init__(self, tp_size: int):
         super().__init__()
@@ -167,10 +227,9 @@ class Flux2AttnProcessorFusedQKV(Flux2AttnProcessor):
     that torch.compile(fullgraph=True) would not trace through.
     """
 
-    def __init__(self, tp_size: int, compile_mode: bool = False):
+    def __init__(self, tp_size: int):
         super().__init__()
         self.tp_size = tp_size
-        self.compile_mode = compile_mode
         self._nki_qkv = None
 
     def _init_nki(self):
@@ -180,19 +239,8 @@ class Flux2AttnProcessorFusedQKV(Flux2AttnProcessor):
         _dir = os.path.dirname(os.path.abspath(__file__))
         if _dir not in sys.path:
             sys.path.insert(0, _dir)
-        from nki_qkv import fused_qkv_op
-
-        if self.compile_mode:
-            # Outer torch.compile(transformer, ...) traces through fused_qkv_op
-            # via FunctionalTensorMode → nki_kernel_wrapper HOP is captured as
-            # an XLA custom call within the full transformer NEFF.
-            self._nki_qkv = fused_qkv_op
-        else:
-            # Eager / lazy-XLA: compile just the QKV function as a standalone NEFF.
-            # dynamic=False: specialize on concrete shapes so dump_config sees plain ints.
-            def _compiled_qkv(x_T, fused_w):
-                return fused_qkv_op(x_T, fused_w)
-            self._nki_qkv = torch.compile(_compiled_qkv, backend="neuron", dynamic=False)
+        from nki_qkv import fused_qkv_kernel
+        self._nki_qkv = fused_qkv_kernel
 
     def _fused_proj(self, attn, q_attr, k_attr, v_attr, cache_attr, x):
         """Build (once) and apply the fused [H, 3*(H//tp)] QKV weight."""
@@ -268,6 +316,95 @@ class Flux2AttnProcessorFusedQKV(Flux2AttnProcessor):
         return hidden_states
 
 
+class Flux2AttnProcessorFusedQKVFlashAttn(Flux2AttnProcessor):
+    """
+    TP-aware double-stream processor: NKI fused QKV projection + NKI flash attention.
+    Combines Flux2AttnProcessorFusedQKV (QKV side) with Flux2AttnProcessorFlashAttn
+    (attention side).  Activated by --fused-qkv --flash-attn together.
+    """
+    def __init__(self, tp_size: int):
+        super().__init__()
+        self.tp_size = tp_size
+        self._nki_qkv = None
+        self._flash_attn = None
+
+    def _init_nki(self):
+        if self._nki_qkv is None:
+            import sys
+            _dir = os.path.dirname(os.path.abspath(__file__))
+            if _dir not in sys.path:
+                sys.path.insert(0, _dir)
+            from nki_qkv import fused_qkv_kernel
+            self._nki_qkv = fused_qkv_kernel
+        if self._flash_attn is None:
+            import sys
+            _dir = os.path.dirname(os.path.abspath(__file__))
+            if _dir not in sys.path:
+                sys.path.insert(0, _dir)
+            from nki_flash_attn import flux2_flash_attn
+            self._flash_attn = flux2_flash_attn
+
+    def _fused_proj(self, attn, q_attr, k_attr, v_attr, cache_attr, x):
+        fused_w = getattr(attn, cache_attr, None)
+        if fused_w is None:
+            fused_w = _build_fused_qkv_weight(
+                getattr(attn, q_attr).weight,
+                getattr(attn, k_attr).weight,
+                getattr(attn, v_attr).weight,
+            ).to(x.device)
+            setattr(attn, cache_attr, fused_w)
+        local_dim = attn.inner_dim // self.tp_size
+        B, S, H = x.shape
+        x_T = x.reshape(B * S, H).T.contiguous()
+        out = self._nki_qkv(x_T, fused_w)
+        return out.reshape(B, S, -1).split(local_dim, dim=-1)
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None,
+                 attention_mask=None, image_rotary_emb=None, **kwargs):
+        self._init_nki()
+        local_heads = attn.heads // self.tp_size
+        head_dim    = attn.head_dim
+
+        query, key, value = self._fused_proj(
+            attn, "to_q", "to_k", "to_v", "_fused_img_qkv_w", hidden_states)
+        query = query.unflatten(-1, (local_heads, head_dim))
+        key   = key.unflatten(-1,   (local_heads, head_dim))
+        value = value.unflatten(-1, (local_heads, head_dim))
+        query = attn.norm_q(query)
+        key   = attn.norm_k(key)
+
+        if attn.added_kv_proj_dim is not None:
+            encoder_query, encoder_key, encoder_value = self._fused_proj(
+                attn, "add_q_proj", "add_k_proj", "add_v_proj",
+                "_fused_txt_qkv_w", encoder_hidden_states)
+            encoder_query = encoder_query.unflatten(-1, (local_heads, head_dim))
+            encoder_key   = encoder_key.unflatten(-1,   (local_heads, head_dim))
+            encoder_value = encoder_value.unflatten(-1, (local_heads, head_dim))
+            encoder_query = attn.norm_added_q(encoder_query)
+            encoder_key   = attn.norm_added_k(encoder_key)
+            query = torch.cat([encoder_query, query], dim=1)
+            key   = torch.cat([encoder_key,   key],   dim=1)
+            value = torch.cat([encoder_value, value], dim=1)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key   = apply_rotary_emb(key,   image_rotary_emb, sequence_dim=1)
+
+        hidden_states = self._flash_attn(query, key, value)
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
+                [encoder_hidden_states.shape[1],
+                 hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+        if encoder_hidden_states is not None:
+            return hidden_states, encoder_hidden_states
+        return hidden_states
+
+
 class Flux2ParallelSelfAttnProcessorTP(Flux2ParallelSelfAttnProcessor):
     def __init__(self, tp_size: int):
         super().__init__()
@@ -303,6 +440,53 @@ class Flux2ParallelSelfAttnProcessorTP(Flux2ParallelSelfAttnProcessor):
         return attn.to_out(hidden_states)
 
 
+class Flux2ParallelSelfAttnProcessorFlashAttn(Flux2ParallelSelfAttnProcessor):
+    """
+    TP-aware single-stream attention processor using a NKI flash attention kernel.
+    """
+    def __init__(self, tp_size: int):
+        super().__init__()
+        self.tp_size = tp_size
+        self._flash_attn = None
+
+    def _init_nki(self):
+        if self._flash_attn is None:
+            import sys, os
+            _dir = os.path.dirname(os.path.abspath(__file__))
+            if _dir not in sys.path:
+                sys.path.insert(0, _dir)
+            from nki_flash_attn import flux2_flash_attn
+            self._flash_attn = flux2_flash_attn
+
+    def __call__(self, attn, hidden_states, attention_mask=None,
+                 image_rotary_emb=None, **kwargs):
+        self._init_nki()
+        local_heads    = attn.heads // self.tp_size
+        head_dim       = attn.head_dim
+        local_inner    = attn.inner_dim // self.tp_size
+        local_mlp_gate = attn.mlp_hidden_dim * attn.mlp_mult_factor // self.tp_size
+
+        hidden_states = attn.to_qkv_mlp_proj(hidden_states)
+        qkv, mlp_hidden_states = torch.split(
+            hidden_states, [3 * local_inner, local_mlp_gate], dim=-1)
+        query, key, value = qkv.chunk(3, dim=-1)
+        query = query.unflatten(-1, (local_heads, head_dim))
+        key   = key.unflatten(-1,   (local_heads, head_dim))
+        value = value.unflatten(-1, (local_heads, head_dim))
+        query = attn.norm_q(query)
+        key   = attn.norm_k(key)
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key   = apply_rotary_emb(key,   image_rotary_emb, sequence_dim=1)
+
+        hidden_states = self._flash_attn(query, key, value)  # (B, S, H_local, D) BSHD
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+        mlp_hidden_states = attn.mlp_act_fn(mlp_hidden_states)
+        hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1)
+        return attn.to_out(hidden_states)
+
+
 # ---------------------------------------------------------------------------
 # Tensor parallelism
 # ---------------------------------------------------------------------------
@@ -372,7 +556,7 @@ def apply_tp_flux2_transformer(
     model: Flux2Transformer2DModel,
     tp_mesh: DeviceMesh,
     fuse_qkv: bool = False,
-    compile_mode: bool = False,
+    flash_attn: bool = False,
 ) -> Flux2Transformer2DModel:
     tp_size = tp_mesh.size()
     double_plan = {
@@ -395,8 +579,36 @@ def apply_tp_flux2_transformer(
         block.ff_context.linear_in.weight.data = _permute_swiglu_for_tp(
             block.ff_context.linear_in.weight.data, tp_size)
         parallelize_module(block, tp_mesh, double_plan)
-        if fuse_qkv:
-            block.attn.set_processor(Flux2AttnProcessorFusedQKV(tp_size, compile_mode=compile_mode))
+        if fuse_qkv and flash_attn:
+            block.attn.set_processor(Flux2AttnProcessorFusedQKVFlashAttn(tp_size))
+            block.attn._fused_img_qkv_w = _build_fused_qkv_weight(
+                block.attn.to_q.weight,
+                block.attn.to_k.weight,
+                block.attn.to_v.weight,
+            )
+            block.attn._fused_txt_qkv_w = _build_fused_qkv_weight(
+                block.attn.add_q_proj.weight,
+                block.attn.add_k_proj.weight,
+                block.attn.add_v_proj.weight,
+            )
+        elif flash_attn:
+            block.attn.set_processor(Flux2AttnProcessorFlashAttn(tp_size))
+        elif fuse_qkv:
+            block.attn.set_processor(Flux2AttnProcessorFusedQKV(tp_size))
+            # Pre-build fused weights now (real DTensors present) so they are
+            # plain local tensors by the time torch.compile traces the forward.
+            # Lazy build inside _fused_proj fails under Dynamo fake-tensor mode
+            # because _to_local() on a fake DTensor returns the full DTensor.
+            block.attn._fused_img_qkv_w = _build_fused_qkv_weight(
+                block.attn.to_q.weight,
+                block.attn.to_k.weight,
+                block.attn.to_v.weight,
+            )
+            block.attn._fused_txt_qkv_w = _build_fused_qkv_weight(
+                block.attn.add_q_proj.weight,
+                block.attn.add_k_proj.weight,
+                block.attn.add_v_proj.weight,
+            )
         else:
             block.attn.set_processor(Flux2AttnProcessorTP(tp_size))
 
@@ -413,30 +625,13 @@ def apply_tp_flux2_transformer(
         attn.to_out.weight.data = _permute_out_for_tp(
             attn.to_out.weight.data, tp_size, inner_dim, mlp_hidden)
         parallelize_module(block, tp_mesh, single_plan)
-        block.attn.set_processor(Flux2ParallelSelfAttnProcessorTP(tp_size))
+        if flash_attn:
+            block.attn.set_processor(Flux2ParallelSelfAttnProcessorFlashAttn(tp_size))
+        else:
+            block.attn.set_processor(Flux2ParallelSelfAttnProcessorTP(tp_size))
     return model
 
 
-def prebuild_fused_qkv_weights(model: Flux2Transformer2DModel) -> None:
-    """Pre-build and cache fused QKV weights on all Flux2AttnProcessorFusedQKV processors.
-
-    Must be called AFTER model.to(device) and BEFORE torch.compile so that
-    _fused_proj finds the pre-built plain tensor (not a DTensor) during tracing.
-    """
-    for block in model.transformer_blocks:
-        proc = block.attn.processor
-        if not isinstance(proc, Flux2AttnProcessorFusedQKV):
-            continue
-        attn = block.attn
-        dev = next(attn.to_q.parameters()).device
-        if not hasattr(attn, "_fused_img_qkv_w") or attn._fused_img_qkv_w is None:
-            attn._fused_img_qkv_w = _build_fused_qkv_weight(
-                attn.to_q.weight, attn.to_k.weight, attn.to_v.weight).to(dev)
-        if not hasattr(attn, "_fused_txt_qkv_w") or attn._fused_txt_qkv_w is None:
-            attn._fused_txt_qkv_w = _build_fused_qkv_weight(
-                attn.add_q_proj.weight, attn.add_k_proj.weight, attn.add_v_proj.weight).to(dev)
-
-
 def apply_tp_text_encoder(model: Qwen3ForCausalLM, tp_mesh: DeviceMesh) -> Qwen3ForCausalLM:
     """Llama-style TP plan for Qwen3ForCausalLM; no weight permutations needed."""
     layer_plan = {