| |
| """Patch sglang-omni s2-pro audio_decoder: SDPA fallback for flash_attn_with_kvcache. |
| |
| FA3 (sgl_kernel.flash_attn) has no sm_120 kernels; the Fast-AR attends over |
| <=11 positions so plain SDPA is numerically fine and fast. Env-gated: |
| FISH_FORCE_SDPA=1 activates the fallback (default keeps upstream FA3 path). |
| """ |
|
|
| from pathlib import Path |
|
|
| P = Path("/opt/work/sglang-omni/sglang_omni/models/fishaudio_s2_pro/" |
| "fish_speech/models/text2semantic/audio_decoder.py") |
| src = P.read_text() |
|
|
| anchor = ''') -> torch.Tensor: |
| return flash_attn_with_kvcache( |
| q=q, |
| k_cache=k_cache, |
| v_cache=v_cache, |
| k=k, |
| v=v, |
| cache_seqlens=cache_seqlens.contiguous() if cache_seqlens is not None else None, |
| causal=causal, |
| num_splits=num_splits, |
| )''' |
|
|
| replacement = ''') -> torch.Tensor: |
| if FISH_FORCE_SDPA: |
| return _sdpa_attn_with_kvcache(q, k_cache, v_cache, k, v, cache_seqlens, causal) |
| return flash_attn_with_kvcache( |
| q=q, |
| k_cache=k_cache, |
| v_cache=v_cache, |
| k=k, |
| v=v, |
| cache_seqlens=cache_seqlens.contiguous() if cache_seqlens is not None else None, |
| causal=causal, |
| num_splits=num_splits, |
| ) |
| |
| |
| def _sdpa_attn_with_kvcache(q, k_cache, v_cache, k, v, cache_seqlens, causal): |
| """Pure-torch drop-in for flash_attn_with_kvcache (sm_120-safe). |
| |
| q: (b, sq, hq, d); k_cache/v_cache: (b, T, hk, d) mutated in place; |
| k/v: (b, sn, hk, d) appended at cache_seqlens; returns (b, sq, hq, d). |
| """ |
| b, sq, hq, d = q.shape |
| hk = k_cache.shape[2] |
| if cache_seqlens is None: |
| cache_seqlens = torch.zeros(b, dtype=torch.int32, device=q.device) |
| cache_seqlens = cache_seqlens.to(q.device) |
| if k is not None: |
| sn = k.shape[1] |
| pos = cache_seqlens.view(b, 1).long() + torch.arange(sn, device=q.device).view(1, sn) |
| bidx = torch.arange(b, device=q.device).view(b, 1).expand(b, sn) |
| k_cache[bidx, pos] = k |
| v_cache[bidx, pos] = v |
| total = cache_seqlens.long() + sn |
| else: |
| total = cache_seqlens.long() |
| max_t = int(total.max()) |
| kk = k_cache[:, :max_t] |
| vv = v_cache[:, :max_t] |
| if hq != hk: |
| rep = hq // hk |
| kk = kk.repeat_interleave(rep, dim=2) |
| vv = vv.repeat_interleave(rep, dim=2) |
| qt = q.transpose(1, 2) # (b, hq, sq, d) |
| kt = kk.transpose(1, 2) # (b, hq, T, d) |
| vt = vv.transpose(1, 2) |
| t_idx = torch.arange(max_t, device=q.device).view(1, 1, 1, max_t) |
| # query j sits at absolute position total-sq+j; causal => attend to <= that |
| limit = (total.view(b, 1, 1, 1) - sq |
| + torch.arange(sq, device=q.device).view(1, 1, sq, 1)) |
| mask = t_idx <= limit # (b, 1, sq, T) broadcast over heads |
| y = torch.nn.functional.scaled_dot_product_attention( |
| qt, kt, vt, attn_mask=mask) |
| return y.transpose(1, 2).contiguous()''' |
|
|
| assert anchor in src, "anchor not found - upstream changed" |
| src = src.replace(anchor, replacement, 1) |
|
|
| env_anchor = 'FISH_BATCH_INVARIANT = os.getenv("FISH_BATCH_INVARIANT", "false").lower() in (' |
| env_add = ('FISH_FORCE_SDPA = os.getenv("FISH_FORCE_SDPA", "false").lower() in ' |
| '("true", "1", "yes")\n') |
| assert env_anchor in src |
| src = src.replace(env_anchor, env_add + env_anchor, 1) |
|
|
| P.write_text(src) |
| print("audio_decoder patched (FISH_FORCE_SDPA gate)") |
| import ast |
| ast.parse(src) |
| print("syntax OK") |
|
|
