File size: 3,458 Bytes
5c2beba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | #!/usr/bin/env python3
"""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")
|