Spaces:
Running on Zero
Running on Zero
| """Masked / cached attention for the MiniMax-H3 blocks. | |
| musubi's H3 Attention runs dense FlashAttention over the whole packed sequence. install() patches its | |
| forward with two extra modes, selected per block call through `active(...)`: | |
| mask : torch flex_attention with a BlockMask over the whole packed sequence (training, reference) | |
| dense: a masks.DenseDF instead of a BlockMask -- the same training mask as FlashAttention-3 prefix calls | |
| kv : the call's rows attend to cached keys/values of earlier rows plus their own (KV-cache sampling, | |
| FlashAttention-3; see KV_ATTENTION); | |
| with store=list the call's own keys/values are appended per layer. A mask may be combined with | |
| kv to fill the cache (chunk 0 together with the condition rows). | |
| The mode is set inside each block call (not around the loop) so that activation-checkpoint | |
| recomputation in backward sees it too. | |
| """ | |
| from __future__ import annotations | |
| import contextlib | |
| import torch | |
| import torch.nn.functional as F | |
| from torch.nn.attention.flex_attention import flex_attention | |
| from musubi_tuner.minimax_h3 import model as h3model | |
| from .layout import CHUNK_SIZES | |
| from .masks import CHUNK_ROWS, DenseDF | |
| _compiled_flex = torch.compile(flex_attention, dynamic=True) | |
| # KV-cache attention over the concatenated past: "fa3" FlashAttention-3 (23% faster per sampled chunk); | |
| # "sdpa" torch SDPA (the earlier behaviour, kept for tests) | |
| KV_ATTENTION = "fa3" | |
| _state = {"mask": None, "kv": None} | |
| _orig_forward = h3model.Attention.forward | |
| _orig_rope = h3model._apply_rope_split_half # kept for tests/test_samplers.py | |
| def _forward(self, hidden_states, rotation_table=None): | |
| mask, kv = _state["mask"], _state["kv"] | |
| if mask is None and kv is None: | |
| return _orig_forward(self, hidden_states, rotation_table) | |
| b, s, _ = hidden_states.shape | |
| q, k, v = self.qkv_proj(hidden_states).split(self.inner_dim, dim=-1) | |
| q = self.q_norm(q.reshape(b, s, self.num_heads, self.head_dim)) | |
| k = self.k_norm(k.reshape(b, s, self.num_heads, self.head_dim)) | |
| v = v.reshape(b, s, self.num_heads, self.head_dim) | |
| if rotation_table is not None: | |
| q = rope(q, rotation_table) | |
| k = rope(k, rotation_table) | |
| dt = hidden_states.dtype if hidden_states.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 | |
| if isinstance(mask, DenseDF): | |
| out = dense_df(q.to(dt), k.to(dt), v.to(dt), mask) | |
| return self.out_proj(out.reshape(b, s, self.inner_dim)) | |
| q, k, v = (t.transpose(1, 2).contiguous().to(dt) for t in (q, k, v)) # [B,H,S,D] | |
| if kv is not None and kv.get("store") is not None: | |
| kv["store"].append((k, v)) | |
| if mask is not None: | |
| out = _compiled_flex(q, k, v, block_mask=mask) | |
| else: | |
| pieces = kv["past"][self._rtcwm_layer] if kv.get("past") is not None else [] | |
| if kv.get("causal_blocks"): # the call's rows are blocks of n rows, block p seeing the | |
| n = kv["causal_blocks"] # past and blocks <= p (the proxy rows of a span) | |
| pk = torch.cat([_key(p) for p in pieces], 2).transpose(1, 2) | |
| pv = torch.cat([p[1] for p in pieces], 2).transpose(1, 2) | |
| qs, ks, vs = (t.transpose(1, 2) for t in (q, k, v)) | |
| out = torch.cat([_fa3(qs[:, i:i + n], torch.cat([pk, ks[:, :i + n]], 1), torch.cat([pv, vs[:, :i + n]], 1)) | |
| for i in range(0, s, n)], 1) | |
| return self.out_proj(out.reshape(b, s, self.inner_dim)) | |
| if pieces: # concatenated per layer: one layer's copy at a time; | |
| k = torch.cat([_key(p) for p in pieces] + [k], 2) # a piece may carry a pending time shift | |
| v = torch.cat([p[1] for p in pieces] + [v], 2) | |
| if KV_ATTENTION == "sdpa": | |
| out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2) | |
| else: | |
| out = _fa3(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)) | |
| return self.out_proj(out.reshape(b, s, self.inner_dim)) | |
| return self.out_proj(out.transpose(1, 2).reshape(b, s, self.inner_dim)) | |
| def _key(piece): | |
| """Cached keys [B,H,S,D] of a piece (k, v) or (k, v, rotation): the rotation (rtcwm.model.sampler._moved) | |
| moves them on the time axis -- RoPE rotations compose -- without storing a moved copy.""" | |
| if len(piece) == 2: | |
| return piece[0] | |
| return rope(piece[0].transpose(1, 2), piece[2]).transpose(1, 2) | |
| def _fa3(q, k, v): | |
| from flash_attn_interface import flash_attn_func | |
| out = flash_attn_func(q, k, v) | |
| return out[0] if isinstance(out, tuple) else out | |
| def dense_df(q, k, v, lay: DenseDF): | |
| """Attention of a packed training span [head | DUV 0..5 | audio | video 0..n-1] under df_mask, as | |
| unmasked FlashAttention-3 calls over key prefixes (see masks.DenseDF). q, k, v [B,S,H,D].""" | |
| h, a = lay.head_rows, lay.audio_rows | |
| dr, cr = lay.duv_rows // len(CHUNK_SIZES), CHUNK_ROWS | |
| d1 = h + lay.duv_rows | |
| v0 = d1 + a | |
| assert q.shape[1] == v0 + lay.n_chunks * cr, (q.shape, lay) | |
| duv = lambda t, p: t[:, h + p * dr:h + (p + 1) * dr] | |
| vid = lambda t, c: t[:, v0 + c * cr:v0 + (c + 1) * cr] | |
| def main(t): | |
| parts = [t[:, :h], t[:, d1:v0]] | |
| for c in range(lay.n_chunks): | |
| parts += [duv(t, c), vid(t, c)] | |
| return torch.cat(parts, 1) | |
| km, vm = main(k), main(v) | |
| kp, vp = (torch.cat([t[:, :h], t[:, d1:v0], t[:, h:d1]], 1) for t in (k, v)) | |
| st = h + a # static keys lead both orders | |
| n = st + (dr + cr if lay.static_sees_first else 0) | |
| o_st = _fa3(torch.cat([q[:, :h], q[:, d1:v0]], 1), km[:, :n], vm[:, :n]) | |
| o_duv = [_fa3(duv(q, p), kp[:, :st + (p + 1) * dr], vp[:, :st + (p + 1) * dr]) for p in range(len(CHUNK_SIZES))] | |
| o_vid = [_fa3(vid(q, c), km[:, :st + (c + 1) * (dr + cr)], vm[:, :st + (c + 1) * (dr + cr)]) | |
| for c in range(lay.n_chunks)] | |
| return torch.cat([o_st[:, :h], *o_duv, o_st[:, h:], *o_vid], 1) | |
| def rope(hidden_states, rotation_table): | |
| """musubi's _apply_rope_split_half written element-wise. The original multiplies every (token, head, | |
| pair) by its 2x2 rotation with a batched matmul -- one tiny cuBLAS gemv per few thousand elements, | |
| ~40% of a sampling step. Same rotation table, accumulated in fp32 like the matmul.""" | |
| pairs = rotation_table.shape[-3] | |
| r = rotation_table.float() # [B,S,1,pairs,2,2] = [[cos, -sin], [sin, cos]] | |
| a, b = hidden_states[..., :pairs].float(), hidden_states[..., pairs:2 * pairs].float() | |
| rotated = torch.cat((r[..., 0, 0] * a + r[..., 0, 1] * b, r[..., 1, 0] * a + r[..., 1, 1] * b), dim=-1) | |
| return torch.cat((rotated.to(hidden_states.dtype), hidden_states[..., 2 * pairs:]), dim=-1) | |
| _orig_mod_scale_shift, _orig_mod_gate = h3model._mod_scale_shift, h3model._mod_gate # kept for tests | |
| _row_index = {} | |
| def _rows(batch_segments, device): | |
| """Per-row modulation index of the (single-sample) segment plan, built once per forward: every block | |
| of a forward call shares the same plan object.""" | |
| key = (id(batch_segments), device) | |
| hit = _row_index.get(key) | |
| if hit is not None and hit[0] is batch_segments: | |
| return hit[1] | |
| (segments,) = batch_segments | |
| seg = torch.tensor(segments, dtype=torch.long) | |
| idx = torch.repeat_interleave(seg[:, 2], seg[:, 1] - seg[:, 0]).to(device, non_blocking=True) | |
| _row_index.clear() | |
| _row_index[key] = (batch_segments, idx) | |
| return idx | |
| def mod_scale_shift(hidden_states, shift, scale, batch_segments): | |
| """musubi's _mod_scale_shift without the Python loop over segments (hundreds of small kernels per block | |
| with the text's token tags): the same x * (1 + scale) + shift, gathered per row.""" | |
| idx = _rows(batch_segments, hidden_states.device) | |
| return hidden_states * (1.0 + scale[idx]).unsqueeze(0) + shift[idx].unsqueeze(0) | |
| def mod_gate(residual, update, gate, batch_segments): | |
| """musubi's _mod_gate, vectorised the same way: residual + update * gate.""" | |
| idx = _rows(batch_segments, residual.device) | |
| return residual + update * gate[idx].unsqueeze(0) | |
| def install(model) -> None: | |
| """Patch H3 attention, its RoPE and its AdaLN modulation (idempotent) and number the layers for the KV | |
| cache.""" | |
| h3model.Attention.forward = _forward | |
| h3model._apply_rope_split_half = rope # the released (unmasked) forward too | |
| h3model._mod_scale_shift, h3model._mod_gate = mod_scale_shift, mod_gate | |
| for i, blk in enumerate(model.blocks): | |
| (blk.attn if hasattr(blk, "attn") else blk.module.attn)._rtcwm_layer = i # FSDP-wrapped or not | |
| def active(mask=None, kv=None): | |
| _state["mask"], _state["kv"] = mask, kv | |
| try: | |
| yield | |
| finally: | |
| _state["mask"], _state["kv"] = None, None | |