初始提交：3 款纯音频游戏 + 共享引擎 + 完整文档与资产管线
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
This commit is contained in:
SoundGame Dev
@@ -0,0 +1,13 @@
|
||||
"""Current-GPU SM arch as a plain 'sm_XX' string.
|
||||
|
||||
Used by swa_mma_aot to retarget the AOT SWA kernel's PTX `.target`. The plain form (no
|
||||
architecture-specific 'a' suffix, e.g. sm_90a / sm_120a) is required because TRT's PTX
|
||||
loader rejects the arch-specific variants; the m16n8k8 TF32 MMA this kernel uses is
|
||||
available from sm_80 up, so the plain target compiles and loads on every arch.
|
||||
"""
|
||||
|
||||
|
||||
def detect_arch() -> str:
|
||||
import torch
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
return f"sm_{major}{minor}"
|
||||
@@ -0,0 +1,252 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Sweep the eight shipped SAME autoencoder configurations: VRAM, latency, accuracy.
|
||||
|
||||
{SAME-S, SAME-L} x {16-bit, fp8} x {chunked, single-shot}
|
||||
|
||||
16-bit means fp16 on SAME-L and bf16 on SAME-S -- the precision each model actually ships.
|
||||
|
||||
python bench_autoencoders.py --out bench.json # full sweep
|
||||
python bench_autoencoders.py --out bench.json --skip-accuracy # ms + VRAM only
|
||||
python make_ae_charts.py bench.json charts.html # render
|
||||
|
||||
── Two things this script exists to get right ────────────────────────────────────────────
|
||||
|
||||
1. ACCURACY IS MEASURED WITH CONTENT HELD FIXED. The obvious construction -- take the
|
||||
first L latents, round-trip them, score against the original -- makes the L axis a
|
||||
content walk, because a longer L is a different, longer piece of music. Measured: at a
|
||||
fixed L, moving the excerpt swings SNR by 9.4-15.0 dB, while the whole L axis at a fixed
|
||||
offset moves 2.95 dB. Content dominates by ~4x, so that chart plots the music, not the
|
||||
model. Instead `--accuracy fixed` (the default) reassembles ONE region from L-sized
|
||||
blocks, so only the processing length changes. That comes out flat to 0.002 dB (SAME-L)
|
||||
above L=16, which is the evidence that the windowed overlap is exact.
|
||||
`--accuracy absolute` keeps the naive construction for the absolute number.
|
||||
|
||||
2. REAL MUSIC, NEVER NOISE. The autoencoder is content-dependent -- published figures for
|
||||
this family span roughly -23 dB (solo piano) to -4 dB (dense circus) -- and noise
|
||||
fabricates numbers that do not hold on music. Point --music at a directory of audio, or
|
||||
at a .npy of shape (2, samples) at 44.1 kHz. L=8192 needs 761 s, longer than most single
|
||||
songs, so several tracks are concatenated rather than one tiled (tiling one track is
|
||||
artificially self-similar).
|
||||
|
||||
Waveform SNR systematically understates a perceptual autoencoder: the same pipeline scores
|
||||
~16 dB on generated audio and ~5 dB on real masters. Compare configurations to each other
|
||||
at a given L; do not read the absolute height as a quality verdict, and do not let it
|
||||
substitute for listening.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse, gc, json, sys, time
|
||||
from pathlib import Path
|
||||
|
||||
SCRIPTS = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(SCRIPTS))
|
||||
|
||||
ENGINES = {
|
||||
("same-l", "16"): dict(enc="same-l/enc_fp16_chunkable.trt",
|
||||
dec="same-l/dec_fp16_chunkable_limiter.trt"),
|
||||
("same-l", "fp8"): dict(enc="same-l/enc_fp8_chunkable.trt",
|
||||
dec="same-l/dec_fp8_chunkable_limiter.trt"),
|
||||
("same-s", "16"): dict(enc="same-s/enc_bf16_chunkable.trt",
|
||||
dec="same-s/dec_bf16_chunkable_limiter.trt"),
|
||||
("same-s", "fp8"): dict(enc="same-s/enc_fp8_chunkable.trt",
|
||||
dec="same-s/dec_fp8_chunkable_limiter.trt"),
|
||||
}
|
||||
DEFAULT_LS = [1, 2, 3, 4, 7, 8, 15, 16, 17, 31, 32, 33, 63, 64, 127, 128, 255, 256, 257,
|
||||
511, 512, 1023, 1024, 1291, 1292, 2047, 2048, 4095, 4096, 4097, 6144, 8192]
|
||||
CHUNK_LAT = 256
|
||||
SPL = 4096 # samples per latent
|
||||
SR = 44100
|
||||
FIXED_REGION = 2048 # latents (~190 s) for the content-held-fixed accuracy pass
|
||||
|
||||
|
||||
def load_music(spec: str, need_samples: int):
|
||||
"""(2, N) float32 @ 44.1 kHz from a .npy, a file, or a directory of audio."""
|
||||
import numpy as np
|
||||
p = Path(spec)
|
||||
if p.suffix == ".npy":
|
||||
a = np.load(p)
|
||||
else:
|
||||
import soundfile as sf
|
||||
from scipy.signal import resample_poly
|
||||
files = sorted(p.glob("**/*")) if p.is_dir() else [p]
|
||||
segs = []
|
||||
got = 0
|
||||
for f in files:
|
||||
if f.suffix.lower() not in (".wav", ".flac", ".mp3", ".ogg", ".m4a"):
|
||||
continue
|
||||
try:
|
||||
x, sr = sf.read(f, dtype="float32", always_2d=True)
|
||||
except Exception:
|
||||
continue
|
||||
x = x.T
|
||||
if x.shape[0] == 1:
|
||||
x = np.repeat(x, 2, 0)
|
||||
x = x[:2]
|
||||
if sr != SR: # a 48 kHz corpus silently fails an == check
|
||||
x = resample_poly(x, SR, sr, axis=1).astype("float32")
|
||||
if x.shape[1] < SR * 5 or float(np.abs(x).max()) < 1e-3:
|
||||
continue
|
||||
segs.append(x); got += x.shape[1]
|
||||
if got >= need_samples:
|
||||
break
|
||||
if not segs:
|
||||
sys.exit(f"error: no usable audio under {spec}")
|
||||
# Concatenate DIFFERENT tracks rather than tile one: a tiled track is
|
||||
# artificially self-similar and flatters the autoencoder.
|
||||
a = np.concatenate(segs, axis=1)
|
||||
if a.shape[1] < need_samples:
|
||||
print(f" warning: only {a.shape[1]/SR:.0f}s of audio; lengths needing more "
|
||||
f"will be skipped")
|
||||
return a.astype("float32")
|
||||
|
||||
|
||||
def snr_db(ref, y):
|
||||
import numpy as np
|
||||
n = min(len(ref), len(y)); ref, y = ref[:n], y[:n]
|
||||
e = ref - y
|
||||
return float(20 * np.log10(np.linalg.norm(ref) / max(np.linalg.norm(e), 1e-12)))
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--out", required=True, help="write results here as JSON")
|
||||
ap.add_argument("--music", default=None,
|
||||
help="directory / file / .npy of REAL music (required unless "
|
||||
"--skip-accuracy)")
|
||||
ap.add_argument("--lengths", default=None,
|
||||
help="comma-separated latent lengths (default: 32 values, 1..8192)")
|
||||
ap.add_argument("--accuracy", choices=("fixed", "absolute", "both"), default="fixed",
|
||||
help="fixed: one region reassembled from L-sized blocks, so only the "
|
||||
"processing length varies (default). absolute: one L-length "
|
||||
"excerpt per L -- the naive construction, dominated by content.")
|
||||
ap.add_argument("--skip-accuracy", action="store_true")
|
||||
ap.add_argument("--repeats", type=int, default=3, help="timing repeats (min is kept)")
|
||||
a = ap.parse_args()
|
||||
|
||||
import sa3_trt_core as canon
|
||||
canon._import_heavy()
|
||||
import numpy as np
|
||||
torch = canon.torch
|
||||
from sa3_trt_core import (TRTRunner, encoder_encode, encode_chunked,
|
||||
decoder_decode, decode_chunked, profile_bounds)
|
||||
A = canon.ARCH_DIR
|
||||
LS = ([int(x) for x in a.lengths.split(",")] if a.lengths else DEFAULT_LS)
|
||||
|
||||
music = None
|
||||
if not a.skip_accuracy:
|
||||
if not a.music:
|
||||
sys.exit("error: --music is required unless --skip-accuracy "
|
||||
"(never benchmark these autoencoders on noise)")
|
||||
music = load_music(a.music, max(LS) * SPL + SR)
|
||||
print(f" music: {music.shape[1]/SR:.0f}s peak {np.abs(music).max():.3f}", flush=True)
|
||||
|
||||
def bind(rel, prof):
|
||||
r = TRTRunner(A / rel, None, True, profile=prof)
|
||||
need = r.engine.get_device_memory_size_for_profile_v2(prof)
|
||||
buf = torch.empty(need, dtype=torch.uint8, device="cuda")
|
||||
r.context.set_device_memory(buf.data_ptr(), need)
|
||||
return r, buf, need
|
||||
|
||||
def fastest(fn):
|
||||
fn(); torch.cuda.synchronize()
|
||||
ts = []
|
||||
for _ in range(a.repeats):
|
||||
torch.cuda.synchronize(); t0 = time.perf_counter(); fn()
|
||||
torch.cuda.synchronize(); ts.append((time.perf_counter() - t0) * 1000)
|
||||
return min(ts)
|
||||
|
||||
rows = []
|
||||
for (model, prec), files in ENGINES.items():
|
||||
ref16 = ENGINES[(model, "16")]
|
||||
for chunked in (True, False):
|
||||
prof = 0 if chunked else 1
|
||||
print(f"\n══ {model} / {prec} / {'chunked' if chunked else 'single-shot'} ══",
|
||||
flush=True)
|
||||
e_r, e_b, e_sc = bind(files["enc"], prof)
|
||||
d_r, d_b, d_sc = bind(files["dec"], prof)
|
||||
# the OTHER stage is pinned to this model's 16-bit chunked engine, so each
|
||||
# accuracy figure is attributable to the stage under test
|
||||
pe_r, pe_b, _ = bind(ref16["enc"], 0)
|
||||
pd_r, pd_b, _ = bind(ref16["dec"], 0)
|
||||
e_hi = profile_bounds(e_r, "audio")[1] // SPL
|
||||
d_hi = profile_bounds(d_r, "latent")[1]
|
||||
|
||||
def roundtrip(off, L, enc_r, enc_ch, enc_hi, dec_r, dec_ch, dec_hi):
|
||||
seg = music[:, off * SPL:(off + L) * SPL]
|
||||
x = torch.from_numpy(np.ascontiguousarray(seg)).unsqueeze(0).cuda()
|
||||
z = (encode_chunked(enc_r, x, chunk_lat=min(CHUNK_LAT, enc_hi))
|
||||
if enc_ch else encoder_encode(enc_r, x))
|
||||
z = z.to(dec_r.in_dtype["latent"])
|
||||
o = (decode_chunked(dec_r, z, chunk_lat=min(CHUNK_LAT, dec_hi))
|
||||
if dec_ch else decoder_decode(dec_r, z))
|
||||
return np.clip(o.float()[0].cpu().numpy() / 32767.0, -1, 1)
|
||||
|
||||
for L in LS:
|
||||
rec = dict(model=model, prec=prec, chunked=chunked, L=L,
|
||||
enc_scratch_mb=e_sc / 2**20, dec_scratch_mb=d_sc / 2**20)
|
||||
# ── latency ──
|
||||
if music is not None and L * SPL <= music.shape[1]:
|
||||
x = torch.from_numpy(
|
||||
np.ascontiguousarray(music[:, :L * SPL])).unsqueeze(0).cuda()
|
||||
else:
|
||||
x = None
|
||||
if x is not None:
|
||||
try:
|
||||
if L > e_hi and not chunked:
|
||||
raise ValueError(f"L>{e_hi} exceeds the single-shot encoder band")
|
||||
ef = ((lambda: encode_chunked(e_r, x, chunk_lat=min(CHUNK_LAT, e_hi)))
|
||||
if chunked else (lambda: encoder_encode(e_r, x)))
|
||||
rec["enc_ms"] = fastest(ef); z = ef()
|
||||
except Exception as ex:
|
||||
rec["enc_err"] = f"{type(ex).__name__}: {ex}"[:70]; z = None
|
||||
try:
|
||||
if z is None:
|
||||
raise ValueError("no latents")
|
||||
if L > d_hi and not chunked:
|
||||
raise ValueError(f"L>{d_hi} exceeds the single-shot decoder band")
|
||||
zz = z.to(d_r.in_dtype["latent"])
|
||||
df = ((lambda: decode_chunked(d_r, zz, chunk_lat=min(CHUNK_LAT, d_hi)))
|
||||
if chunked else (lambda: decoder_decode(d_r, zz)))
|
||||
rec["dec_ms"] = fastest(df)
|
||||
except Exception as ex:
|
||||
rec["dec_err"] = f"{type(ex).__name__}: {ex}"[:70]
|
||||
# ── accuracy ──
|
||||
stages = {"enc": (e_r, chunked, e_hi, pd_r, True, CHUNK_LAT),
|
||||
"dec": (pe_r, True, CHUNK_LAT, d_r, chunked, d_hi)}
|
||||
if music is not None and a.accuracy in ("fixed", "both") \
|
||||
and L <= FIXED_REGION and FIXED_REGION % L == 0:
|
||||
ref = music[:, :FIXED_REGION * SPL].T
|
||||
for which, args_ in stages.items():
|
||||
try:
|
||||
if (not args_[1] and L > args_[2]) or (not args_[4] and L > args_[5]):
|
||||
raise ValueError("band ceiling")
|
||||
y = np.concatenate([roundtrip(o, L, *args_)
|
||||
for o in range(0, FIXED_REGION, L)], axis=0)
|
||||
rec[f"{which}_fix_db"] = snr_db(ref, y)
|
||||
except Exception as ex:
|
||||
rec[f"{which}_fix_err"] = str(ex)[:50]
|
||||
if music is not None and a.accuracy in ("absolute", "both") \
|
||||
and L * SPL <= music.shape[1]:
|
||||
ref = music[:, :L * SPL].T
|
||||
for which, args_ in stages.items():
|
||||
try:
|
||||
if (not args_[1] and L > args_[2]) or (not args_[4] and L > args_[5]):
|
||||
raise ValueError("band ceiling")
|
||||
rec[f"{which}_rt_db"] = snr_db(ref, roundtrip(0, L, *args_))
|
||||
except Exception as ex:
|
||||
rec[f"{which}_rt_err"] = str(ex)[:50]
|
||||
rows.append(rec)
|
||||
print(f" L={L:>5d} enc {rec.get('enc_ms', float('nan')):8.2f}ms "
|
||||
f"dec {rec.get('dec_ms', float('nan')):8.2f}ms "
|
||||
f"fixed {rec.get('dec_fix_db', float('nan')):7.3f}dB", flush=True)
|
||||
json.dump(rows, open(a.out, "w"), indent=1)
|
||||
for o in (e_r, d_r, pe_r, pd_r):
|
||||
del o
|
||||
del e_b, d_b, pe_b, pd_b
|
||||
gc.collect(); torch.cuda.empty_cache()
|
||||
print(f"\n wrote {a.out} ({len(rows)} rows)")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,229 @@
|
||||
"""DiT-only timing benchmark across multiple engines + L values.
|
||||
|
||||
Loads each engine once, runs dit.step() in a tight loop at every requested L
|
||||
(skipping L below the engine's profile min), captures fp32 timings with
|
||||
torch.cuda.Event, and reports median across N runs.
|
||||
|
||||
Usage:
|
||||
python bench_dit_profile.py \
|
||||
--engines '<label>=<path>,<label>=<path>,...' \
|
||||
--lvals 32,64,128,256,512,1024,2048,4096 \
|
||||
--warmup 3 --runs 5 [--seed 0]
|
||||
|
||||
Engines run sequentially; each is loaded, benchmarked at all valid L values,
|
||||
then freed before the next engine. Results are emitted as a TSV table to
|
||||
stdout (engine, L, median_ms, mean_ms, min_ms, max_ms, runs).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse
|
||||
import contextlib
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
SCRIPTS = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(SCRIPTS))
|
||||
|
||||
|
||||
# ── lazy heavy imports (torch + trt + plugin) ─────────────────────────────
|
||||
@contextlib.contextmanager
|
||||
def _silence_fd(fileno: int):
|
||||
saved = os.dup(fileno)
|
||||
try:
|
||||
with open(os.devnull, "wb") as null:
|
||||
os.dup2(null.fileno(), fileno)
|
||||
yield
|
||||
finally:
|
||||
os.dup2(saved, fileno)
|
||||
os.close(saved)
|
||||
|
||||
|
||||
def _import_heavy():
|
||||
"""Mirror sa3_trt._import_heavy() so its module globals (torch, trt) are
|
||||
populated — TRTRunner/DiTRunner reach into sa3_trt's globals.
|
||||
"""
|
||||
global torch, trt
|
||||
import torch as _torch
|
||||
import tensorrt as _trt
|
||||
torch = _torch
|
||||
trt = _trt
|
||||
# Inject torch/trt into sa3_trt_core's globals (formerly sa3_trt before the
|
||||
# full-pipeline-graph refactor); its classes TRTRunner / DiTRunner reach
|
||||
# into the module globals to find these.
|
||||
import sa3_trt_core as _s
|
||||
_s.torch = _torch
|
||||
_s.trt = _trt
|
||||
with _silence_fd(1), _silence_fd(2):
|
||||
try:
|
||||
import diff_attn_nocast_plugin # noqa: F401
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _engine_profile_min_L(engine, ctx) -> int:
|
||||
"""Read input-shape lower bound on 'x' from the engine's optimization profile.
|
||||
|
||||
TRT exposes profile shapes via Engine.get_tensor_profile_shape(name, profile_idx).
|
||||
Returns the min L on the last (sequence) axis of 'x'.
|
||||
"""
|
||||
lo, opt, hi = engine.get_tensor_profile_shape("x", 0)
|
||||
return int(lo[-1])
|
||||
|
||||
|
||||
def _engine_profile_max_L(engine) -> int:
|
||||
lo, opt, hi = engine.get_tensor_profile_shape("x", 0)
|
||||
return int(hi[-1])
|
||||
|
||||
|
||||
def bench_one_engine(engine_path: str, label: str, lvals: list[int],
|
||||
warmup: int, runs: int, seed: int) -> list[dict]:
|
||||
"""Load `engine_path`, benchmark dit.step() at each L in `lvals` that fits
|
||||
within the engine's profile. Returns one dict per (L) measurement.
|
||||
"""
|
||||
from sa3_trt_core import TRTRunner, DiTRunner, IO_CHANNELS, T5_MAX_LEN, COND_DIM
|
||||
|
||||
print(f"\n━━━ engine: {label} ━━━", flush=True)
|
||||
print(f" path: {engine_path}", flush=True)
|
||||
sz = Path(engine_path).stat().st_size
|
||||
print(f" size: {sz/1e9:.2f} GB", flush=True)
|
||||
|
||||
runner = TRTRunner(Path(engine_path))
|
||||
dit = DiTRunner(runner)
|
||||
|
||||
lo = _engine_profile_min_L(runner.engine, runner.context)
|
||||
hi = _engine_profile_max_L(runner.engine)
|
||||
print(f" profile min/max L: {lo}/{hi}", flush=True)
|
||||
|
||||
torch.manual_seed(seed)
|
||||
g = torch.Generator(device="cuda")
|
||||
g.manual_seed(seed)
|
||||
|
||||
rows = []
|
||||
valid_lvals = [L for L in lvals if lo <= L <= hi]
|
||||
skipped = [L for L in lvals if L < lo or L > hi]
|
||||
if skipped:
|
||||
print(f" skipping L outside profile: {skipped}", flush=True)
|
||||
|
||||
for L in valid_lvals:
|
||||
# Inputs sized for this L (fp32 throughout; DiTRunner casts internally
|
||||
# in step()).
|
||||
x = torch.randn(1, IO_CHANNELS, L, device="cuda", dtype=torch.float32, generator=g)
|
||||
t_val = torch.tensor([0.5], device="cuda", dtype=torch.float32)
|
||||
t5_h = torch.randn(1, T5_MAX_LEN, COND_DIM, device="cuda", dtype=torch.float32, generator=g)
|
||||
t5_m = torch.ones(1, T5_MAX_LEN, device="cuda", dtype=torch.float32)
|
||||
sec = torch.tensor([float(L * 4096 / 44100)], device="cuda", dtype=torch.float32)
|
||||
lac = torch.zeros(1, 257, L, device="cuda", dtype=torch.float32)
|
||||
|
||||
# Warmup (also forces _setup at this L, which sets input shape on the ctx).
|
||||
for _ in range(warmup):
|
||||
_ = dit.step(x, t_val, t5_h, t5_m, sec, lac)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Timed runs (per-step timing — the engine's stream is sync'd inside
|
||||
# step() already, so each iteration is one DiT forward end-to-end).
|
||||
times_ms = []
|
||||
for _ in range(runs):
|
||||
ev0 = torch.cuda.Event(enable_timing=True)
|
||||
ev1 = torch.cuda.Event(enable_timing=True)
|
||||
ev0.record()
|
||||
_ = dit.step(x, t_val, t5_h, t5_m, sec, lac)
|
||||
ev1.record()
|
||||
torch.cuda.synchronize()
|
||||
times_ms.append(ev0.elapsed_time(ev1))
|
||||
|
||||
med = float(np.median(times_ms))
|
||||
mn = float(np.mean(times_ms))
|
||||
lo_t = float(min(times_ms))
|
||||
hi_t = float(max(times_ms))
|
||||
print(f" L={L:>5} median={med:7.2f} ms mean={mn:7.2f} min={lo_t:7.2f} max={hi_t:7.2f} "
|
||||
f"(n={runs})", flush=True)
|
||||
rows.append({
|
||||
"engine": label,
|
||||
"engine_path": engine_path,
|
||||
"engine_size_bytes": sz,
|
||||
"profile_min_L": lo,
|
||||
"profile_max_L": hi,
|
||||
"L": L,
|
||||
"median_ms": med,
|
||||
"mean_ms": mn,
|
||||
"min_ms": lo_t,
|
||||
"max_ms": hi_t,
|
||||
"runs": runs,
|
||||
})
|
||||
|
||||
# Free the inputs before the next L (they grow with L for x/lac).
|
||||
del x, t5_h, t5_m, t_val, sec, lac
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Free the engine + context before loading the next one.
|
||||
runner.free()
|
||||
del runner, dit
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return rows
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--engines", required=True,
|
||||
help="Comma-sep list of label=path entries, e.g. "
|
||||
"'canonical=/.../dit_bf16.trt,min32=/.../dit_bf16_min32.trt'")
|
||||
ap.add_argument("--lvals", default="32,64,128,256,512,1024,2048,4096",
|
||||
help="Comma-sep L values to benchmark at. Engines skip L < their profile min.")
|
||||
ap.add_argument("--warmup", type=int, default=3)
|
||||
ap.add_argument("--runs", type=int, default=5)
|
||||
ap.add_argument("--seed", type=int, default=0)
|
||||
ap.add_argument("--tsv", default=None,
|
||||
help="Optional path to write rows as TSV (in addition to stdout).")
|
||||
args = ap.parse_args()
|
||||
|
||||
_import_heavy()
|
||||
|
||||
engines = []
|
||||
for chunk in args.engines.split(","):
|
||||
if not chunk.strip():
|
||||
continue
|
||||
if "=" not in chunk:
|
||||
sys.exit(f"--engines entry must be 'label=path', got {chunk!r}")
|
||||
label, path = chunk.split("=", 1)
|
||||
engines.append((label.strip(), path.strip()))
|
||||
|
||||
lvals = [int(x) for x in args.lvals.split(",") if x.strip()]
|
||||
|
||||
all_rows = []
|
||||
for label, path in engines:
|
||||
rows = bench_one_engine(path, label, lvals, args.warmup, args.runs, args.seed)
|
||||
all_rows.extend(rows)
|
||||
|
||||
# Final summary table — pivot by L → engine.
|
||||
print("\n━━━ summary (median ms, n={}) ━━━".format(args.runs))
|
||||
eng_labels = [e[0] for e in engines]
|
||||
header = ["L"] + eng_labels
|
||||
print("\t".join(header))
|
||||
by_eng_L = {(r["engine"], r["L"]): r for r in all_rows}
|
||||
for L in lvals:
|
||||
row = [str(L)]
|
||||
for el in eng_labels:
|
||||
r = by_eng_L.get((el, L))
|
||||
row.append(f"{r['median_ms']:.2f}" if r else "—")
|
||||
print("\t".join(row))
|
||||
|
||||
if args.tsv:
|
||||
import csv
|
||||
keys = ["engine", "engine_path", "engine_size_bytes", "profile_min_L",
|
||||
"profile_max_L", "L", "median_ms", "mean_ms", "min_ms", "max_ms", "runs"]
|
||||
with open(args.tsv, "w") as f:
|
||||
w = csv.DictWriter(f, fieldnames=keys, delimiter="\t")
|
||||
w.writeheader()
|
||||
for r in all_rows:
|
||||
w.writerow(r)
|
||||
print(f"\nwrote {args.tsv} ({len(all_rows)} rows)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,143 @@
|
||||
"""
|
||||
TRT plugin for diff SWA attention — NO dtype cast, minimal overhead.
|
||||
|
||||
Two implementations of the SAME op are registered:
|
||||
|
||||
aot_impl (preferred) An ahead-of-time PTX kernel compiled INTO the engine at
|
||||
build time. No Python, no Triton, nothing host-side in the
|
||||
enqueue path — so the engine is CUDA-graph-capturable.
|
||||
impl (fallback) The original Triton path, which re-enters Python on every
|
||||
enqueue. TRT uses this only if the AOT path is unavailable.
|
||||
|
||||
Why the AOT path matters: with `impl` alone the built engine is NOT stream-capturable
|
||||
on sm_120. `enqueueV3` returns False under capture, TRT reports "this TRT engine is
|
||||
not stream capturable" from myelin/runner.cpp, and the stage is silently omitted from
|
||||
the graph — the SAME-L decode vanished from the mega-graph and every render returned
|
||||
the pre-capture warmup decode of zero latents (a constant wash, identical across
|
||||
seeds and prompts, exit code 0). It happened to capture on sm_90, but re-entering
|
||||
Python inside a captured region is fragile by construction rather than supported.
|
||||
|
||||
Capturability is baked into the engine at BUILD time, so engines must be rebuilt to
|
||||
pick this up. The ONNX is unchanged — same op name, same signature.
|
||||
|
||||
The AOT kernel's contract, verified against what TRT actually passes this op in the
|
||||
SAME-L decoder: FP32 in/out, q/k/v each (B, N, 2H, D) = (1, 4352, 48, 64), D=64,
|
||||
2H=48 so H=24, window=17. It performs the differential subtraction in FP32 inside
|
||||
the kernel, which is also what the Triton path does via o[:, :, :H] - o[:, :, H:].
|
||||
"""
|
||||
import os
|
||||
|
||||
import torch
|
||||
import tensorrt.plugin as trtp
|
||||
from typing import Tuple
|
||||
|
||||
from diff_swa_ptx_kernel import generate_diff_swa_ptx
|
||||
|
||||
WINDOW = 17
|
||||
HEAD_DIM = 64
|
||||
WARPS_PER_BLOCK = 4
|
||||
_ptx_cache: dict = {}
|
||||
|
||||
# Which AOT kernel to compile into the engine:
|
||||
# "mma" block-tiled Triton-compiled kernel with tensor cores (swa_mma_aot).
|
||||
# BLOCK_N queries per block sharing the K/V window in SRAM; ~3.6x faster per
|
||||
# layer than "ptx" and on par with the JIT path.
|
||||
# "ptx" the original hand-written scalar-FP32 kernel (diff_swa_ptx_kernel): one warp
|
||||
# per query, no K/V reuse, zero mma instructions.
|
||||
# Both are graph-capturable; this only trades speed. Override with SA3_SWA_AOT.
|
||||
SWA_AOT_BACKEND = os.environ.get("SA3_SWA_AOT", "mma").lower()
|
||||
|
||||
|
||||
def _ptx_for(num_heads: int):
|
||||
"""(kernel_name, ptx) for this head count, generated once and cached."""
|
||||
if num_heads not in _ptx_cache:
|
||||
_ptx_cache[num_heads] = generate_diff_swa_ptx(
|
||||
window=WINDOW, D=HEAD_DIM, H=num_heads,
|
||||
warps_per_block=WARPS_PER_BLOCK)
|
||||
return _ptx_cache[num_heads]
|
||||
|
||||
_stream_cache = {}
|
||||
_triton_fn = None
|
||||
|
||||
|
||||
@trtp.register("samel::diff_attn_swa")
|
||||
def diff_attn_swa_desc(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
|
||||
v_bat: trtp.TensorDesc, num_heads: int) -> trtp.TensorDesc:
|
||||
out = q_bat.like()
|
||||
out.shape_expr[-2] = q_bat.shape_expr[-2] // 2
|
||||
return out
|
||||
|
||||
|
||||
@trtp.impl("samel::diff_attn_swa")
|
||||
def diff_attn_swa_impl(q_bat: trtp.Tensor, k_bat: trtp.Tensor, v_bat: trtp.Tensor,
|
||||
num_heads: int, outputs: Tuple[trtp.Tensor], stream: int):
|
||||
global _triton_fn
|
||||
if stream not in _stream_cache:
|
||||
_stream_cache[stream] = torch.cuda.ExternalStream(stream)
|
||||
if _triton_fn is None:
|
||||
from triton_swa_v2 import triton_swa_attn_v2
|
||||
_triton_fn = triton_swa_attn_v2
|
||||
|
||||
with torch.cuda.stream(_stream_cache[stream]):
|
||||
q = torch.as_tensor(q_bat, device="cuda")
|
||||
k = torch.as_tensor(k_bat, device="cuda")
|
||||
v = torch.as_tensor(v_bat, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
# NO dtype cast — Triton auto-compiles for the input dtype
|
||||
o = _triton_fn(q, k, v, window=17)
|
||||
H = num_heads
|
||||
out_t.copy_(o[:, :, :H, :] - o[:, :, H:, :])
|
||||
|
||||
|
||||
@trtp.aot_impl("samel::diff_attn_swa")
|
||||
def diff_attn_swa_aot(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
|
||||
v_bat: trtp.TensorDesc, num_heads: int,
|
||||
outputs: Tuple[trtp.TensorDesc], tactic: int = None
|
||||
) -> Tuple[str, str, trtp.KernelLaunchParams, trtp.SymIntExprs]:
|
||||
"""AOT PTX variant — makes the engine graph-capturable. See module docstring.
|
||||
|
||||
Grid is (ceil(N / warps_per_block), H_out, B); each warp handles one query
|
||||
position for head h and head h+H, subtracting in FP32 before the store. Strides
|
||||
are passed as symbolic extras so the kernel works across the dynamic N axis.
|
||||
"""
|
||||
B = q_bat.shape_expr[0]
|
||||
N = q_bat.shape_expr[1]
|
||||
H2 = q_bat.shape_expr[2]
|
||||
D = q_bat.shape_expr[3]
|
||||
H_out = H2 // 2
|
||||
|
||||
if SWA_AOT_BACKEND == "mma":
|
||||
# Block-tiled tensor-core kernel. Its scalars are baked in as constexprs at
|
||||
# compile time, so the only runtime extra is N. Grid is one block per
|
||||
# (query tile, output head, batch); shared memory holds the Q/K/V tiles.
|
||||
import swa_mma_aot
|
||||
from _arch import detect_arch
|
||||
name, ptx, shared, warps = swa_mma_aot.build(num_heads, detect_arch())
|
||||
extra = trtp.SymIntExprs(1)
|
||||
extra[0] = N
|
||||
return (name, ptx,
|
||||
trtp.KernelLaunchParams(
|
||||
grid_x=trtp.cdiv(N, swa_mma_aot.BLOCK_N),
|
||||
grid_y=H_out,
|
||||
grid_z=B,
|
||||
block_x=warps * 32,
|
||||
shared_mem=shared),
|
||||
extra)
|
||||
|
||||
# Hand-written scalar kernel: one warp per query, strides passed at runtime.
|
||||
name, ptx = _ptx_for(num_heads)
|
||||
extra = trtp.SymIntExprs(5)
|
||||
extra[0] = N
|
||||
extra[1] = H2 * D # stride between positions, input
|
||||
extra[2] = D # stride between heads, input
|
||||
extra[3] = H_out * D # stride between positions, output
|
||||
extra[4] = D # stride between heads, output
|
||||
|
||||
return (name, ptx,
|
||||
trtp.KernelLaunchParams(
|
||||
grid_x=trtp.cdiv(N, WARPS_PER_BLOCK),
|
||||
grid_y=H_out,
|
||||
grid_z=B,
|
||||
block_x=WARPS_PER_BLOCK * 32),
|
||||
extra)
|
||||
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
Generate PTX for DIFFERENTIAL SWA attention — subtraction built into the kernel.
|
||||
Avoids FP16 catastrophic cancellation by computing (attn_primary - attn_diff) in FP32.
|
||||
|
||||
Grid: (ceil(N/warps_per_block), H, B) — H is num_heads (24), NOT 2*H
|
||||
Each warp: computes SWA for head h (primary) and head h+H (diff), then subtracts.
|
||||
Output: (B, N, H, D) — already-subtracted differential result.
|
||||
|
||||
Plugin signature: diff_swa_attn(Q_bat, K_bat, V_bat) where Q_bat is (B, N, 2H, D)
|
||||
Output: (B, N, H, D)
|
||||
"""
|
||||
|
||||
def generate_diff_swa_ptx(window=17, D=64, H=24, warps_per_block=4):
|
||||
"""Generate PTX for differential SWA. Returns (kernel_name, ptx_string)."""
|
||||
WIN = 2 * window + 1
|
||||
H2 = 2 * H # Total heads in input (48)
|
||||
THREADS_PER_BLOCK = warps_per_block * 32
|
||||
elem_shift = 2 # log2(4 bytes for fp32)
|
||||
kernel_name = "diff_swa_attn_ptx"
|
||||
|
||||
ptx = f"""
|
||||
.version 8.0
|
||||
.target sm_90
|
||||
.address_size 64
|
||||
|
||||
.visible .entry {kernel_name}(
|
||||
.param .u64 .ptr .global .align 4 Q_ptr,
|
||||
.param .u64 .ptr .global .align 4 K_ptr,
|
||||
.param .u64 .ptr .global .align 4 V_ptr,
|
||||
.param .u32 N_param,
|
||||
.param .u32 stride_n_in,
|
||||
.param .u32 stride_h_in,
|
||||
.param .u32 stride_n_out,
|
||||
.param .u32 stride_h_out,
|
||||
.param .u64 .ptr .global .align 4 Out_ptr
|
||||
) {{
|
||||
.reg .pred %p<16>;
|
||||
.reg .b32 %r<80>;
|
||||
.reg .f32 %f<160>;
|
||||
.reg .b64 %rd<40>;
|
||||
|
||||
// Thread indexing: {warps_per_block} warps per block
|
||||
mov.u32 %r3, %tid.x;
|
||||
shr.u32 %r40, %r3, 5; // warp_id
|
||||
and.b32 %r41, %r3, 31; // lane_id
|
||||
|
||||
// Position: block_n * {warps_per_block} + warp_id
|
||||
mov.u32 %r42, %ctaid.x;
|
||||
mov.u32 %r43, {warps_per_block};
|
||||
mul.lo.u32 %r44, %r42, %r43;
|
||||
add.u32 %r0, %r44, %r40; // n = position
|
||||
|
||||
mov.u32 %r1, %ctaid.y; // h = output head (0..{H-1})
|
||||
mov.u32 %r2, %ctaid.z; // b = batch
|
||||
|
||||
ld.param.u32 %r4, [N_param];
|
||||
ld.param.u32 %r5, [stride_n_in]; // input stride_n (for 2H heads)
|
||||
ld.param.u32 %r7, [stride_h_in]; // input stride_h
|
||||
ld.param.u32 %r50, [stride_n_out]; // output stride_n (for H heads)
|
||||
ld.param.u32 %r51, [stride_h_out]; // output stride_h
|
||||
|
||||
setp.ge.u32 %p0, %r0, %r4;
|
||||
@%p0 bra DONE;
|
||||
|
||||
// Compute INPUT base offset for PRIMARY head h: n*stride_n + h*stride_h + lane*2
|
||||
mul.lo.u32 %r6, %r0, %r5;
|
||||
mad.lo.u32 %r8, %r1, %r7, %r6;
|
||||
shl.b32 %r9, %r41, 1;
|
||||
add.u32 %r10, %r8, %r9;
|
||||
|
||||
// Compute INPUT base offset for DIFF head h+{H}: n*stride_n + (h+{H})*stride_h + lane*2
|
||||
mov.u32 %r52, {H};
|
||||
add.u32 %r53, %r1, %r52; // h + H
|
||||
mad.lo.u32 %r54, %r53, %r7, %r6; // n*stride_n + (h+H)*stride_h
|
||||
add.u32 %r55, %r54, %r9; // + lane*2
|
||||
|
||||
// Load Q for PRIMARY head
|
||||
cvt.u64.u32 %rd0, %r10;
|
||||
shl.b64 %rd1, %rd0, {elem_shift};
|
||||
ld.param.u64 %rd2, [Q_ptr];
|
||||
add.u64 %rd3, %rd2, %rd1;
|
||||
ld.global.f32 %f0, [%rd3];
|
||||
ld.global.f32 %f1, [%rd3 + 4];
|
||||
|
||||
// Load Q for DIFF head
|
||||
cvt.u64.u32 %rd20, %r55;
|
||||
shl.b64 %rd21, %rd20, {elem_shift};
|
||||
add.u64 %rd22, %rd2, %rd21;
|
||||
ld.global.f32 %f2, [%rd22];
|
||||
ld.global.f32 %f3, [%rd22 + 4];
|
||||
|
||||
// Scale
|
||||
mov.f32 %f14, 0f3E000000; // 0.125 = 1/sqrt(64)
|
||||
mov.f32 %f40, 0f3FB8AA3B; // log2(e)
|
||||
|
||||
// Online softmax for PRIMARY
|
||||
mov.f32 %f10, 0fFF800000; mov.f32 %f11, 0f00000000;
|
||||
mov.f32 %f12, 0f00000000; mov.f32 %f13, 0f00000000;
|
||||
// Online softmax for DIFF
|
||||
mov.f32 %f60, 0fFF800000; mov.f32 %f61, 0f00000000;
|
||||
mov.f32 %f62, 0f00000000; mov.f32 %f63, 0f00000000;
|
||||
"""
|
||||
|
||||
for wi in range(WIN):
|
||||
ptx += f"""
|
||||
// --- Window position {wi} ---
|
||||
mov.s32 %r21, %r0;
|
||||
add.s32 %r21, %r21, {wi - window};
|
||||
setp.lt.s32 %p1, %r21, 0;
|
||||
setp.ge.s32 %p2, %r21, %r4;
|
||||
or.pred %p3, %p1, %p2;
|
||||
@%p3 bra SKIP_{wi};
|
||||
|
||||
// K offset for kp: kp*stride_n + h*stride_h + lane*2 (PRIMARY)
|
||||
mul.lo.u32 %r22, %r21, %r5;
|
||||
mad.lo.u32 %r23, %r1, %r7, %r22;
|
||||
add.u32 %r24, %r23, %r9;
|
||||
cvt.u64.u32 %rd4, %r24;
|
||||
shl.b64 %rd5, %rd4, {elem_shift};
|
||||
ld.param.u64 %rd6, [K_ptr];
|
||||
add.u64 %rd7, %rd6, %rd5;
|
||||
ld.global.f32 %f20, [%rd7];
|
||||
ld.global.f32 %f21, [%rd7 + 4];
|
||||
|
||||
// K offset for DIFF head: kp*stride_n + (h+H)*stride_h + lane*2
|
||||
mad.lo.u32 %r56, %r53, %r7, %r22;
|
||||
add.u32 %r57, %r56, %r9;
|
||||
cvt.u64.u32 %rd23, %r57;
|
||||
shl.b64 %rd24, %rd23, {elem_shift};
|
||||
add.u64 %rd25, %rd6, %rd24;
|
||||
ld.global.f32 %f70, [%rd25];
|
||||
ld.global.f32 %f71, [%rd25 + 4];
|
||||
|
||||
// PRIMARY dot product
|
||||
mul.f32 %f22, %f0, %f20;
|
||||
fma.rn.f32 %f22, %f1, %f21, %f22;
|
||||
mov.b32 %r30, %f22;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 16, 0x1f, 0xffffffff; mov.b32 %f23, %r30; add.f32 %f22, %f22, %f23;
|
||||
mov.b32 %r30, %f22;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 8, 0x1f, 0xffffffff; mov.b32 %f23, %r30; add.f32 %f22, %f22, %f23;
|
||||
mov.b32 %r30, %f22;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 4, 0x1f, 0xffffffff; mov.b32 %f23, %r30; add.f32 %f22, %f22, %f23;
|
||||
mov.b32 %r30, %f22;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 2, 0x1f, 0xffffffff; mov.b32 %f23, %r30; add.f32 %f22, %f22, %f23;
|
||||
mov.b32 %r30, %f22;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 1, 0x1f, 0xffffffff; mov.b32 %f23, %r30; add.f32 %f22, %f22, %f23;
|
||||
mul.f32 %f22, %f22, %f14;
|
||||
|
||||
// DIFF dot product
|
||||
mul.f32 %f72, %f2, %f70;
|
||||
fma.rn.f32 %f72, %f3, %f71, %f72;
|
||||
mov.b32 %r30, %f72;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 16, 0x1f, 0xffffffff; mov.b32 %f73, %r30; add.f32 %f72, %f72, %f73;
|
||||
mov.b32 %r30, %f72;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 8, 0x1f, 0xffffffff; mov.b32 %f73, %r30; add.f32 %f72, %f72, %f73;
|
||||
mov.b32 %r30, %f72;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 4, 0x1f, 0xffffffff; mov.b32 %f73, %r30; add.f32 %f72, %f72, %f73;
|
||||
mov.b32 %r30, %f72;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 2, 0x1f, 0xffffffff; mov.b32 %f73, %r30; add.f32 %f72, %f72, %f73;
|
||||
mov.b32 %r30, %f72;
|
||||
shfl.sync.bfly.b32 %r30, %r30, 1, 0x1f, 0xffffffff; mov.b32 %f73, %r30; add.f32 %f72, %f72, %f73;
|
||||
mul.f32 %f72, %f72, %f14;
|
||||
|
||||
// PRIMARY online softmax update
|
||||
max.f32 %f30, %f10, %f22;
|
||||
sub.f32 %f32, %f10, %f30; mul.f32 %f32, %f32, %f40; ex2.approx.f32 %f32, %f32;
|
||||
sub.f32 %f33, %f22, %f30; mul.f32 %f33, %f33, %f40; ex2.approx.f32 %f33, %f33;
|
||||
mul.f32 %f11, %f11, %f32; add.f32 %f11, %f11, %f33;
|
||||
mul.f32 %f12, %f12, %f32; mul.f32 %f13, %f13, %f32;
|
||||
|
||||
// DIFF online softmax update
|
||||
max.f32 %f80, %f60, %f72;
|
||||
sub.f32 %f82, %f60, %f80; mul.f32 %f82, %f82, %f40; ex2.approx.f32 %f82, %f82;
|
||||
sub.f32 %f83, %f72, %f80; mul.f32 %f83, %f83, %f40; ex2.approx.f32 %f83, %f83;
|
||||
mul.f32 %f61, %f61, %f82; add.f32 %f61, %f61, %f83;
|
||||
mul.f32 %f62, %f62, %f82; mul.f32 %f63, %f63, %f82;
|
||||
|
||||
// Load V (shared between primary and diff - V[:,:,h,:] == V[:,:,h+H,:])
|
||||
ld.param.u64 %rd8, [V_ptr];
|
||||
add.u64 %rd9, %rd8, %rd5;
|
||||
ld.global.f32 %f34, [%rd9];
|
||||
ld.global.f32 %f35, [%rd9 + 4];
|
||||
|
||||
// PRIMARY accumulate
|
||||
fma.rn.f32 %f12, %f33, %f34, %f12;
|
||||
fma.rn.f32 %f13, %f33, %f35, %f13;
|
||||
|
||||
// DIFF accumulate
|
||||
fma.rn.f32 %f62, %f83, %f34, %f62;
|
||||
fma.rn.f32 %f63, %f83, %f35, %f63;
|
||||
|
||||
mov.f32 %f10, %f30;
|
||||
mov.f32 %f60, %f80;
|
||||
|
||||
SKIP_{wi}:
|
||||
"""
|
||||
|
||||
ptx += f"""
|
||||
// Normalize and SUBTRACT in FP32
|
||||
rcp.approx.f32 %f50, %f11;
|
||||
mul.f32 %f12, %f12, %f50;
|
||||
mul.f32 %f13, %f13, %f50;
|
||||
|
||||
rcp.approx.f32 %f90, %f61;
|
||||
mul.f32 %f62, %f62, %f90;
|
||||
mul.f32 %f63, %f63, %f90;
|
||||
|
||||
// DIFFERENTIAL: primary - diff (in FP32!)
|
||||
sub.f32 %f12, %f12, %f62;
|
||||
sub.f32 %f13, %f13, %f63;
|
||||
|
||||
// Store to OUTPUT: (B, N, H, D) layout using output strides
|
||||
mul.lo.u32 %r60, %r0, %r50; // n * stride_n_out
|
||||
mad.lo.u32 %r61, %r1, %r51, %r60; // + h * stride_h_out
|
||||
add.u32 %r62, %r61, %r9; // + lane*2
|
||||
cvt.u64.u32 %rd10, %r62;
|
||||
shl.b64 %rd11, %rd10, {elem_shift};
|
||||
ld.param.u64 %rd12, [Out_ptr];
|
||||
add.u64 %rd13, %rd12, %rd11;
|
||||
st.global.f32 [%rd13], %f12;
|
||||
st.global.f32 [%rd13 + 4], %f13;
|
||||
|
||||
DONE:
|
||||
ret;
|
||||
}}
|
||||
"""
|
||||
return kernel_name, ptx
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
name, ptx = generate_diff_swa_ptx()
|
||||
print(f"Kernel: {name}, PTX: {len(ptx)} bytes, shared: {'.shared' in ptx}")
|
||||
import subprocess
|
||||
with open('/tmp/diff_swa.ptx', 'w') as f: f.write(ptx)
|
||||
r = subprocess.run(['ptxas', '--gpu-name=sm_90', '-o', '/tmp/diff_swa.cubin', '/tmp/diff_swa.ptx'],
|
||||
capture_output=True, text=True)
|
||||
print(f"ptxas: {'OK' if r.returncode == 0 else 'FAILED'}")
|
||||
if r.returncode != 0: print(r.stderr[:300])
|
||||
@@ -0,0 +1,310 @@
|
||||
"""Backend-agnostic gradio UI helpers, copied from optimized/mlx/scripts/sa3_gradio.py.
|
||||
|
||||
Forked rather than shared: the MLX and TensorRT trees install independently, with their own
|
||||
requirements, and a common module would couple two deployment units that today have no reason to
|
||||
know about each other. Nothing in here touches a backend -- these are player/history/queue
|
||||
renderers and filename helpers -- so the two copies only need to be reconciled when the UI
|
||||
itself changes.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import html as html_lib
|
||||
import re
|
||||
import time
|
||||
import urllib.parse
|
||||
import wave
|
||||
|
||||
|
||||
SAMPLE_RATE = 44100 # every SA3 codec is 44.1 kHz stereo
|
||||
|
||||
_JS_FIX_SLIDERS = (
|
||||
"() => setTimeout(() => {"
|
||||
"document.querySelectorAll('input[type=range]').forEach(r => {"
|
||||
"const min = +r.min || 0, max = +r.max || 100, v = +r.value;"
|
||||
"if (max > min) r.style.setProperty('--range_progress',"
|
||||
"((v - min) / (max - min) * 100) + '%');"
|
||||
"});}, 150)"
|
||||
)
|
||||
|
||||
_JS_PAUSE_OTHERS = (
|
||||
"if(this.isConnected){"
|
||||
"var t=this;var L=window._sa3All=(window._sa3All||[]);"
|
||||
"if(L.indexOf(t)<0)L.push(t);"
|
||||
"L.forEach(function(o){if(o!==t){try{o.pause()}catch(e){}}});"
|
||||
"document.querySelectorAll('audio').forEach(function(o){if(o!==t)o.pause();});"
|
||||
"window._sa3All=L.filter(function(o){return o===t||o.isConnected;});}"
|
||||
"else{this.pause();}")
|
||||
|
||||
_JS_PLAYHEAD = ("var p=this.closest('.blk').querySelector('.ph');"
|
||||
"if(p&&this.duration)p.style.left=(this.currentTime/this.duration*100)+'%';")
|
||||
# Main-player position ledger (for Hotswap): every timeupdate/play/pause stamps
|
||||
# the current position so a freshly swapped-in element can resume exactly there.
|
||||
# Guards (all verified via CDP against the live page):
|
||||
# - isConnected: teardown fires a pause AND a final timeupdate on the removed
|
||||
# element — detached elements never write the ledger.
|
||||
# - unresumed hotswap candidates (data-hs set, hsd not yet) don't write either:
|
||||
# with autoplay, the NEW element's 'play' event fires BEFORE loadedmetadata
|
||||
# and would stamp t≈0 over the old clip's position right before the resume
|
||||
# handler reads it.
|
||||
|
||||
_JS_POS_RECORD = ("if(this.isConnected&&!(this.dataset.hs==='1'&&!this.dataset.hsd))"
|
||||
"{window._sa3Pos={t:this.currentTime,playing:!this.paused,ts:Date.now()};}")
|
||||
# timeupdate variant: a PAUSED element only fires timeupdate on seeks — e.g. the
|
||||
# resume handler's own currentTime assignment. Stamping playing:false there
|
||||
# poisons the ledger for any second render of the same clip (gradio sometimes
|
||||
# renders a component update twice), which froze the handoff chain. While
|
||||
# paused, only real pause events may write.
|
||||
|
||||
_JS_POS_RECORD_TU = ("if(this.isConnected&&!this.paused&&"
|
||||
"!(this.dataset.hs==='1'&&!this.dataset.hsd))"
|
||||
"{window._sa3Pos={t:this.currentTime,playing:true,ts:Date.now()};}")
|
||||
# Resilient play: Chrome's autoplay policy can reject programmatic play() once
|
||||
# the transient user activation from the Generate click has expired (seconds),
|
||||
# leaving the clip correctly positioned but paused. Try, retry shortly after,
|
||||
# and as a last resort arm a one-shot listener so the user's next click or
|
||||
# keypress anywhere resumes playback.
|
||||
# Every attempt re-checks isConnected: a superseded (detached) element must
|
||||
# never revive itself from a queued retry and fight the current player.
|
||||
|
||||
_JS_TRY_PLAY = (
|
||||
"var A=this;A.play().catch(function(){setTimeout(function(){"
|
||||
"if(!A.isConnected)return;"
|
||||
"A.play().catch(function(){console.warn('play blocked by autoplay policy — "
|
||||
"will resume on next interaction');"
|
||||
"var f=function(){if(A.isConnected)A.play();"
|
||||
"document.removeEventListener('pointerdown',f,true);"
|
||||
"document.removeEventListener('keydown',f,true);};"
|
||||
"document.addEventListener('pointerdown',f,true);"
|
||||
"document.addEventListener('keydown',f,true);});},150);});")
|
||||
# Hotswap resume: if the previous main audio was playing when this one arrived,
|
||||
# jump to its position (+ the split-second since the last stamp) and keep going;
|
||||
# beyond the new clip's duration -> start at zero. Guarded (hsd) so it applies
|
||||
# once, on whichever of loadedmetadata/canplay fires first.
|
||||
|
||||
_JS_HOTSWAP = ("if(this.isConnected&&this.dataset.hs==='1'&&!this.dataset.hsd&&this.duration){"
|
||||
"this.dataset.hsd=1;var s=window._sa3Pos;"
|
||||
"this.dataset.dbg=s?(s.playing?'playing@'+s.t.toFixed(2):'notplaying@'+s.t.toFixed(2)):'noledger';"
|
||||
"if(s&&s.playing){var tt=s.t+(Date.now()-s.ts)/1000;"
|
||||
"this.currentTime=(tt<this.duration)?tt:0;" + _JS_TRY_PLAY + "}}")
|
||||
_JS_SEEK = ("var a=this.closest('.blk').querySelector('audio');"
|
||||
"var r=this.getBoundingClientRect();"
|
||||
"if(a&&a.duration){a.currentTime=(event.clientX-r.left)/r.width*a.duration;a.play();}")
|
||||
|
||||
_JS_PROMOTE = ("var b=document.getElementById('sa3-promote');"
|
||||
"if(b){(b.tagName==='BUTTON'?b:b.querySelector('button')||b).click();}")
|
||||
# Scroll anchoring for the history panel: onscroll continuously records which
|
||||
# item sits at the viewport top (+offset); after every re-render a hidden
|
||||
# bootstrap <img onerror> restores that anchor — stick-to-top when at top,
|
||||
# otherwise keep hovering over the same old item as new ones prepend.
|
||||
|
||||
_JS_SCROLL_RECORD = (
|
||||
"var c=this,k=null,off=0,ch=c.querySelectorAll('[data-key]');"
|
||||
"for(var i=0;i<ch.length;i++){if(ch[i].offsetTop+ch[i].offsetHeight>c.scrollTop)"
|
||||
"{k=ch[i].getAttribute('data-key');off=ch[i].offsetTop-c.scrollTop;break}}"
|
||||
"window._sa3S={atTop:c.scrollTop<8,key:k,off:off};")
|
||||
|
||||
_JS_SCROLL_RESTORE = (
|
||||
"var c=document.getElementById('sa3-hist');var s=window._sa3S||{};"
|
||||
"if(c){var el=s.key?c.querySelector('[data-key="'+s.key+'"]'):null;"
|
||||
"if(el&&!s.atTop){c.scrollTop=el.offsetTop-(s.off||0)}else{c.scrollTop=0}}"
|
||||
"this.remove();")
|
||||
|
||||
|
||||
def _ago(ts) -> str:
|
||||
if not ts:
|
||||
return ""
|
||||
d = max(0.0, time.time() - ts)
|
||||
if d < 10:
|
||||
return "just now"
|
||||
if d < 60:
|
||||
return f"{int(d)}s ago"
|
||||
if d < 3600:
|
||||
return f"{int(d // 60)}m ago"
|
||||
if d < 86400:
|
||||
return f"{int(d // 3600)}h ago"
|
||||
return f"{int(d // 86400)}d ago"
|
||||
|
||||
def condense_prompt(prompt: str) -> str:
|
||||
"""Prompt → filename fragment (the main repo's verbose-naming rule):
|
||||
filesystem-special characters become hyphens, capped at 150 chars."""
|
||||
prompt = re.sub(r'[\\/:*?"<>|]', '-', prompt)[:150]
|
||||
return prompt or "_"
|
||||
|
||||
def verbose_basename(prompt, negative_prompt, cfg, sigma_max, seed) -> str:
|
||||
"""prompt[.neg-…].cfg{scale}[.smx{σ}].{seed} — matches the main repo's
|
||||
gradio 'verbose' file naming (cfg segment only when cfg != 1)."""
|
||||
base = condense_prompt(prompt)
|
||||
if negative_prompt and negative_prompt.strip():
|
||||
base += ".neg-" + condense_prompt(negative_prompt.strip())
|
||||
if cfg != 1.0:
|
||||
base += f".cfg{cfg:g}"
|
||||
if sigma_max != 1.0:
|
||||
base += f".smx{sigma_max:g}"
|
||||
return f"{base}.{seed}"
|
||||
|
||||
def _save_wav(pcm_int16, out_path):
|
||||
"""pcm_int16: (T, 2) int16 interleaved."""
|
||||
with wave.open(str(out_path), "wb") as w:
|
||||
w.setnchannels(2)
|
||||
w.setsampwidth(2)
|
||||
w.setframerate(SAMPLE_RATE)
|
||||
w.writeframes(pcm_int16.tobytes())
|
||||
|
||||
|
||||
def _meta_suffix(entry) -> str:
|
||||
"""' · cfg 1.5 · noise 0.92 · audio2audio · inpainting' — only the non-defaults.
|
||||
The sigma tag reads 'noise' for a2a runs (it IS the init_noise_level there),
|
||||
'smx' otherwise."""
|
||||
parts = []
|
||||
cfg = entry.get("cfg", 1.0)
|
||||
smx = entry.get("smx", 1.0)
|
||||
mode = entry.get("mode", "")
|
||||
is_a2a = "a2a" in mode or mode == "audio-to-audio"
|
||||
dit = entry.get("dit", "")
|
||||
if dit:
|
||||
parts.append({"medium": "med", "sm-music": "sm-mus", "sm-sfx": "sm-sfx"}.get(dit, dit))
|
||||
if cfg != 1.0:
|
||||
parts.append(f"cfg {cfg:g}")
|
||||
if smx != 1.0:
|
||||
parts.append(f"{'noise' if is_a2a else 'smx'} {smx:g}")
|
||||
if is_a2a:
|
||||
parts.append("audio2audio")
|
||||
if "inpaint" in mode:
|
||||
parts.append("inpainting")
|
||||
if entry.get("lora"):
|
||||
parts.append(f"lora {entry['lora']}")
|
||||
return "".join(f" · {p}" for p in parts)
|
||||
|
||||
def _neg_disp(entry) -> str:
|
||||
"""Labeled negative prompt, shown only when it actually acted (cfg != 1)."""
|
||||
neg = (entry.get("neg") or "").strip()
|
||||
if neg and entry.get("cfg", 1.0) != 1.0:
|
||||
return f' · <span style="opacity:0.75">neg: {html_lib.escape(neg)}</span>'
|
||||
return ""
|
||||
|
||||
def render_player(entry, *, small=False, autoplay=False, autodl=False, radio=False,
|
||||
bg=None, advance=False, loop=False, hotswap=False):
|
||||
"""One self-contained player block: audio + caption + seekable spectrogram
|
||||
with playhead. Global one-at-a-time playback via onplay pause-others.
|
||||
small: audio + spectrogram side by side (half width each), 'Xm ago' caption.
|
||||
advance: on ended, hop to the next history item's audio (Auto-play).
|
||||
loop: native loop attribute — finished audio restarts (suppresses ended).
|
||||
hotswap: main slot only — resume at the previous clip's position on arrival."""
|
||||
# assemble per-event handler bodies (a duplicated attribute name would
|
||||
# silently drop one handler, so each event is emitted exactly once)
|
||||
on_canplay = []
|
||||
if small:
|
||||
attrs = f'onplay="{_JS_PAUSE_OTHERS}" ontimeupdate="{_JS_PLAYHEAD}"'
|
||||
else:
|
||||
attrs = (f'onplay="{_JS_PAUSE_OTHERS}{_JS_POS_RECORD}" '
|
||||
f'onpause="{_JS_POS_RECORD}" '
|
||||
f'ontimeupdate="{_JS_PLAYHEAD}{_JS_POS_RECORD_TU}"')
|
||||
if hotswap and not small:
|
||||
attrs += f' data-hs="1" onloadedmetadata="{_JS_HOTSWAP}"'
|
||||
on_canplay.append(_JS_HOTSWAP) # fallback if loadedmetadata already passed
|
||||
if autoplay and not small:
|
||||
# the autoplay attribute is subject to the same policy — rescue a
|
||||
# policy-paused clip (radio transitions, long generations)
|
||||
on_canplay.append("if(this.paused&&!(this.dataset.hs==='1'&&!this.dataset.hsd))"
|
||||
"{" + _JS_TRY_PLAY + "}")
|
||||
if autodl:
|
||||
js_name = entry["name"].replace("\\", "").replace("'", "\\'")
|
||||
on_canplay.append("if(!this.dataset.dld){this.dataset.dld=1;"
|
||||
"var l=document.createElement('a');l.href=this.src;"
|
||||
f"l.download='{js_name}';l.click();}}")
|
||||
if on_canplay:
|
||||
attrs += f' oncanplay="{"".join(on_canplay)}"'
|
||||
if loop:
|
||||
attrs += " loop"
|
||||
elif radio:
|
||||
attrs += f' onended="{_JS_PROMOTE}"'
|
||||
elif advance:
|
||||
attrs += (' onended="var b=this.closest(\'.blk\');'
|
||||
"var n=b?b.nextElementSibling:null;"
|
||||
"while(n&&!n.classList.contains('blk'))n=n.nextElementSibling;"
|
||||
"if(n){var a=n.querySelector('audio');if(a)a.play();}\"")
|
||||
auto = "autoplay " if autoplay else ""
|
||||
# Serve audio via gradio's file route instead of a data: URI — the URL ends
|
||||
# with the real (verbose) filename, so right-click "Save audio as…" offers
|
||||
# it instead of download.mp3, and the page stays light as history grows.
|
||||
src = "gradio_api/file=" + urllib.parse.quote(entry["path"], safe="/")
|
||||
audio_el = (f'<audio controls {auto}style="width:100%" {attrs} '
|
||||
f'src="{src}"></audio>')
|
||||
prompt_disp = html_lib.escape(entry["prompt"]) or "<i>(no prompt)</i>"
|
||||
|
||||
spec_core = ""
|
||||
if entry.get("spec_b64"):
|
||||
height = "height:56px;" if small else ""
|
||||
tip = ("3-band tinted stereo mel · red=bass / green=mid / blue=high · "
|
||||
"L top, R bottom · click to seek")
|
||||
spec_core = (f'<div style="position:relative; cursor:pointer" title="{tip}" '
|
||||
f'onclick="{_JS_SEEK}">'
|
||||
f'<img src="data:image/png;base64,{entry["spec_b64"]}" '
|
||||
f'style="width:100%; {height} display:block; image-rendering:pixelated; '
|
||||
f'border:1px solid #333" alt="spectrogram"/>'
|
||||
f'<div class="ph" style="position:absolute; top:0; bottom:0; left:0%; '
|
||||
f'width:2px; background:#fff; pointer-events:none; '
|
||||
f'box-shadow:0 0 4px rgba(0,0,0,.8)"></div>'
|
||||
f'</div>')
|
||||
|
||||
if small:
|
||||
row = (f'<div style="display:flex; gap:8px; align-items:center">'
|
||||
f'<div style="flex:1; min-width:0">{audio_el}</div>'
|
||||
f'<div style="flex:1; min-width:0">{spec_core}</div></div>')
|
||||
cap = (f'<div style="font-size:0.8em; margin:2px 0; color:#888">'
|
||||
f'{_ago(entry.get("ts"))} · {prompt_disp}{_neg_disp(entry)} · seed {entry["seed"]}'
|
||||
f'{_meta_suffix(entry)}</div>')
|
||||
style = f"padding:6px 8px;{' background:' + bg + ';' if bg else ''}"
|
||||
return (f'<div class="blk" data-key="{entry["key"]}" style="{style}">'
|
||||
f'{row}{cap}</div>')
|
||||
|
||||
return (f'<div class="blk" data-key="{entry["key"]}">'
|
||||
f'{audio_el}{spec_core}</div>')
|
||||
|
||||
def render_history(hist, advance=False, loop=False):
|
||||
if not hist:
|
||||
return ""
|
||||
# zebra striping instead of separators: light grey vs medium grey rows
|
||||
items = "".join(
|
||||
render_player(e, small=True, advance=advance, loop=loop,
|
||||
bg="rgba(127,127,127,0.24)" if i % 2 else "rgba(127,127,127,0.08)")
|
||||
for i, e in enumerate(hist))
|
||||
boot = f'<img src="data:," style="display:none" onerror="{_JS_SCROLL_RESTORE}"/>'
|
||||
return (f'<div style="font-weight:600; margin-top:14px">Previous generations ({len(hist)})</div>'
|
||||
f'<div id="sa3-hist" onscroll="{_JS_SCROLL_RECORD}" '
|
||||
f'style="max-height:480px; overflow-y:auto; position:relative; margin-top:6px; '
|
||||
f'padding-right:6px">{boot}{items}</div>')
|
||||
|
||||
def render_queue_status(entry=None, generating=False):
|
||||
"""Status chip for the pre-rendered next take.
|
||||
|
||||
Once a take is ready the chip is a CONTROL, not a label: clicking it promotes the
|
||||
queued render and plays it immediately. It fires the same hidden #sa3-promote button
|
||||
radio's onended uses, so both routes go through one handler — and because the click
|
||||
is a real user gesture, the browser's autoplay policy lets the audio start.
|
||||
|
||||
While a take is still rendering there is nothing to promote, so the chip stays inert.
|
||||
"""
|
||||
if generating:
|
||||
body, ready = "generating…", False
|
||||
elif entry is not None:
|
||||
p = html_lib.escape(entry["prompt"]) or "<i>(no prompt)</i>"
|
||||
body = f"ready — {p}{_neg_disp(entry)} · seed {entry['seed']}{_meta_suffix(entry)}"
|
||||
ready = True
|
||||
else:
|
||||
return ""
|
||||
style = ("margin-top:2px; padding:6px 10px; background:rgba(127,127,127,0.12); "
|
||||
"border-radius:6px; color:#888; font-size:0.85em")
|
||||
if not ready:
|
||||
return f'<div style="{style}"><b>Queued next</b> · {body}</div>'
|
||||
# Enter/Space as well as click, so it behaves like the button it now is.
|
||||
keys = ("if(event.key===\'Enter\'||event.key===\' \'){event.preventDefault();"
|
||||
+ _JS_PROMOTE + "}")
|
||||
return (f'<div class="sa3-queued" role="button" tabindex="0" '
|
||||
f'title="Play this take now" '
|
||||
f'style="{style}; cursor:pointer; user-select:none" '
|
||||
f'onclick="{_JS_PROMOTE}" onkeydown="{keys}">'
|
||||
f'<b>Queued next</b> · {body}'
|
||||
f'<span style="float:right; opacity:0.75">▶ play now</span></div>')
|
||||
|
||||
|
||||
# ── Gradio UI ──────────────────────────────────────────────────────────────
|
||||
@@ -0,0 +1,194 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Render bench_autoencoders.py output into one self-contained HTML page.
|
||||
|
||||
python make_ae_charts.py bench.json charts.html
|
||||
|
||||
Six charts -- VRAM, latency, accuracy, for the encoder and the decoder -- eight lines
|
||||
each. No external assets: the SVG is inline, so the page works offline.
|
||||
"""
|
||||
import json, sys, math, html
|
||||
|
||||
if len(sys.argv) < 3:
|
||||
raise SystemExit("usage: make_ae_charts.py <bench.json> <out.html>")
|
||||
rows = json.load(open(sys.argv[1]))
|
||||
OUT = sys.argv[2]
|
||||
|
||||
SERIES = [ # (model, prec, chunked) -> label, colour, dash
|
||||
("same-l", "16", True ), ("same-l", "16", False),
|
||||
("same-l", "fp8", True ), ("same-l", "fp8", False),
|
||||
("same-s", "16", True ), ("same-s", "16", False),
|
||||
("same-s", "fp8", True ), ("same-s", "fp8", False),
|
||||
]
|
||||
COLOR = {("same-l","16"):"#e07b39", ("same-l","fp8"):"#a8341b",
|
||||
("same-s","16"):"#3b9ad9", ("same-s","fp8"):"#1d5b86"}
|
||||
def label(m,p,c):
|
||||
pn = {"16":"16-bit","fp8":"fp8"}[p]
|
||||
return f"{m.upper()} · {pn} · {'chunked' if c else 'single-shot'}"
|
||||
|
||||
def series_points(m,p,c,key):
|
||||
out=[]
|
||||
for r in rows:
|
||||
if (r["model"],r["prec"],r["chunked"])==(m,p,c) and key in r:
|
||||
v=r[key]
|
||||
if isinstance(v,(int,float)) and v==v and v>0:
|
||||
out.append((r["L"], v))
|
||||
return sorted(out)
|
||||
|
||||
W,H = 1180, 560
|
||||
PAD = dict(l=92, r=28, t=26, b=64)
|
||||
def chart(key, title, ylab, logy, unit, note=""):
|
||||
data = {s: series_points(*s, key) for s in SERIES}
|
||||
pts = [v for s in data.values() for _,v in s]
|
||||
xs = sorted({x for s in data.values() for x,_ in s})
|
||||
if not pts: return f"<p>no data for {key}</p>"
|
||||
lo, hi = min(pts), max(pts)
|
||||
if logy:
|
||||
lo = max(lo, hi/1e4); ylo, yhi = math.log10(lo*0.85), math.log10(hi*1.18)
|
||||
ytr = lambda v: math.log10(max(v,lo*0.85))
|
||||
else:
|
||||
span=hi-lo or 1; ylo, yhi = lo-span*0.10, hi+span*0.12; ytr = lambda v: v
|
||||
xlo, xhi = math.log10(min(xs)*0.85), math.log10(max(xs)*1.18)
|
||||
iw, ih = W-PAD['l']-PAD['r'], H-PAD['t']-PAD['b']
|
||||
X = lambda L: PAD['l'] + (math.log10(L)-xlo)/(xhi-xlo)*iw
|
||||
Y = lambda v: PAD['t'] + ih - (ytr(v)-ylo)/(yhi-ylo)*ih
|
||||
g=[f'<svg viewBox="0 0 {W} {H}" class="chart" role="img" aria-label="{html.escape(title)}">']
|
||||
# y grid
|
||||
if logy:
|
||||
ticks=[]; e=math.floor(ylo)
|
||||
while e<=math.ceil(yhi):
|
||||
for mlt in (1,2,5):
|
||||
v=mlt*10**e
|
||||
if lo*0.85<=v<=hi*1.18: ticks.append(v)
|
||||
e+=1
|
||||
else:
|
||||
ticks=[ylo+(yhi-ylo)*i/6 for i in range(7)]
|
||||
for v in ticks:
|
||||
y=Y(v)
|
||||
g.append(f'<line x1="{PAD["l"]}" x2="{W-PAD["r"]}" y1="{y:.1f}" y2="{y:.1f}" class="grid"/>')
|
||||
t=(f"{v:,.0f}" if v>=100 or (not logy and abs(v)>=10) else f"{v:,.2f}".rstrip("0").rstrip("."))
|
||||
g.append(f'<text x="{PAD["l"]-10}" y="{y+4:.1f}" class="ytick">{t}</text>')
|
||||
for L in [1,4,16,64,256,1024,4096,8192]:
|
||||
if not (min(xs)*0.85<=L<=max(xs)*1.18): continue
|
||||
x=X(L)
|
||||
g.append(f'<line y1="{PAD["t"]}" y2="{H-PAD["b"]}" x1="{x:.1f}" x2="{x:.1f}" class="grid"/>')
|
||||
g.append(f'<text x="{x:.1f}" y="{H-PAD["b"]+20}" class="xtick">{L}</text>')
|
||||
g.append(f'<text x="{W/2}" y="{H-PAD["b"]+46}" class="axlab">latent length L '
|
||||
f' (L×4096 samples ≈ L×0.093 s)</text>')
|
||||
g.append(f'<text transform="translate(22,{PAD["t"]+ih/2}) rotate(-90)" class="axlab">'
|
||||
f'{html.escape(ylab)}</text>')
|
||||
for s in SERIES:
|
||||
p=data[s]
|
||||
if not p: continue
|
||||
d=" ".join(("M" if i==0 else "L")+f"{X(L):.1f},{Y(v):.1f}" for i,(L,v) in enumerate(p))
|
||||
col=COLOR[(s[0],s[1])]; dash='6 4' if s[2] else 'none'
|
||||
g.append(f'<path d="{d}" fill="none" stroke="{col}" stroke-width="2.4" '
|
||||
f'stroke-dasharray="{dash}" stroke-linejoin="round"/>')
|
||||
for L,v in p:
|
||||
g.append(f'<circle cx="{X(L):.1f}" cy="{Y(v):.1f}" r="2.6" fill="{col}">'
|
||||
f'<title>{label(*s)} — L={L}: {v:,.2f} {unit}</title></circle>')
|
||||
g.append("</svg>")
|
||||
leg="".join(
|
||||
f'<span class="key"><svg width="30" height="10"><line x1="1" y1="5" x2="29" y2="5" '
|
||||
f'stroke="{COLOR[(s[0],s[1])]}" stroke-width="2.6" '
|
||||
f'stroke-dasharray="{"6 4" if s[2] else "none"}"/></svg>{html.escape(label(*s))}</span>'
|
||||
for s in SERIES)
|
||||
return (f'<section><h2>{html.escape(title)}</h2>'
|
||||
+ (f'<p class="note">{note}</p>' if note else "")
|
||||
+ "".join(g) + f'<div class="legend">{leg}</div></section>')
|
||||
|
||||
CHARTS = [
|
||||
("enc", "enc_scratch_mb", "Encoder — VRAM scratch reserved", "scratch (MB, log)", True, "MB",
|
||||
"TensorRT commits a context's scratch from its PROFILE CEILING, not from the shape you bind — "
|
||||
"so these are flat in L. That is the whole point of the low band. SAME-L's two precisions "
|
||||
"reserve identical scratch, so those lines overlap exactly."),
|
||||
("enc", "enc_ms", "Encoder — latency", "milliseconds (log)", True, "ms",
|
||||
"Best of three, after warm-up. Single-shot lines stop at L=4096: that is the wide band's "
|
||||
"ceiling. Chunked keeps going because it windows."),
|
||||
("enc", "enc_fix_db", "Encoder — accuracy vs the original music, content held fixed",
|
||||
"reconstruction SNR (dB)", False, "dB",
|
||||
"The SAME 190 s of real music at every point — reassembled from L-sized blocks, so only the "
|
||||
"processing length changes. This is the apples-to-apples question the TFLite rung table asks, "
|
||||
"and the answer is the same: flat. From L=16 up the whole curve moves 0.002 dB on SAME-L and "
|
||||
"0.020 dB on SAME-S. The fall below L≈16 is real but is not a seam: 0.09 s blocks encoded "
|
||||
"independently lose context at every boundary."),
|
||||
("dec", "dec_scratch_mb", "Decoder — VRAM scratch reserved", "scratch (MB, log)", True, "MB", ""),
|
||||
("dec", "dec_ms", "Decoder — latency", "milliseconds (log)", True, "ms", ""),
|
||||
("dec", "dec_fix_db", "Decoder — accuracy vs the original music, content held fixed",
|
||||
"reconstruction SNR (dB)", False, "dB",
|
||||
"Same construction, decoder varying. Chunked and single-shot land on the same value to three "
|
||||
"decimals at every length, which is the direct evidence that the windowed overlap is exact."),
|
||||
]
|
||||
body = "".join(chart(k, t, y, ly, u, n) for _, k, t, y, ly, u, n in CHARTS)
|
||||
n_ok = sum(1 for r in rows if "enc_ms" in r)
|
||||
n_len = len({r['L'] for r in rows})
|
||||
CSS = """
|
||||
:root{--bg:#faf8f5;--fg:#1c1a17;--mut:#6b645c;--line:#e0dbd3;--card:#fff;--accent:#a8341b}
|
||||
:root:not([data-theme="light"]){}
|
||||
@media (prefers-color-scheme: dark){:root:not([data-theme="light"]){
|
||||
--bg:#14120f;--fg:#ece7df;--mut:#9a9086;--line:#2e2a25;--card:#1b1815;--accent:#e07b39}}
|
||||
:root[data-theme="dark"]{--bg:#14120f;--fg:#ece7df;--mut:#9a9086;--line:#2e2a25;
|
||||
--card:#1b1815;--accent:#e07b39}
|
||||
*{box-sizing:border-box}
|
||||
body{margin:0;padding:38px 22px 90px;background:var(--bg);color:var(--fg);
|
||||
font:15px/1.6 ui-sans-serif,-apple-system,"Segoe UI",Roboto,sans-serif}
|
||||
main{max-width:1240px;margin:0 auto}
|
||||
h1{font-size:1.75rem;margin:0 0 6px;letter-spacing:-.02em;text-wrap:balance}
|
||||
.sub{color:var(--mut);margin:0 0 30px;max-width:78ch}
|
||||
section{background:var(--card);border:1px solid var(--line);border-radius:10px;
|
||||
padding:20px 22px 16px;margin:0 0 26px;overflow-x:auto}
|
||||
h2{font-size:1.06rem;margin:0 0 4px;letter-spacing:-.01em}
|
||||
.note{color:var(--mut);font-size:.86rem;margin:0 0 12px;max-width:92ch}
|
||||
.chart{width:100%;height:auto;min-width:820px;display:block}
|
||||
.grid{stroke:var(--line);stroke-width:1}
|
||||
.ytick,.xtick{fill:var(--mut);font-size:12px;font-variant-numeric:tabular-nums}
|
||||
.ytick{text-anchor:end}.xtick{text-anchor:middle}
|
||||
.axlab{fill:var(--mut);font-size:12.5px;text-anchor:middle}
|
||||
.legend{display:flex;flex-wrap:wrap;gap:6px 18px;margin-top:12px;
|
||||
border-top:1px solid var(--line);padding-top:12px}
|
||||
.key{display:flex;align-items:center;gap:7px;font-size:.83rem;color:var(--mut)}
|
||||
footer{color:var(--mut);font-size:.85rem;max-width:88ch;margin-top:34px}
|
||||
code{background:rgba(128,128,128,.14);padding:1px 5px;border-radius:4px;font-size:.9em}
|
||||
"""
|
||||
open(OUT, "w").write(f"""<meta charset="utf-8">
|
||||
<title>SAME Autoencoder Sweep</title>
|
||||
<style>{CSS}</style>
|
||||
<main>
|
||||
<h1>SAME autoencoder: VRAM, latency, accuracy</h1>
|
||||
<p class="sub">Eight configurations — {{SAME-S, SAME-L}} × {{16-bit, fp8}} × {{chunked,
|
||||
single-shot}} — measured across {n_len} latent lengths from
|
||||
L=1 to L=8192 ({n_ok} points). Accuracy is a full round trip against <em>real music</em> — seven
|
||||
mastered tracks from the 806k-track corpus (rock, jazz, classical, electronic, hip hop, piano,
|
||||
folk) concatenated to 762 s, because L=8192 is longer than any single song and tiling one
|
||||
track would be artificially self-similar. Never noise: the autoencoder is content-dependent and
|
||||
noise fabricates numbers that do not hold on music. 16-bit means fp16 on SAME-L and bf16 on
|
||||
SAME-S. Hover any point for its value.</p>
|
||||
{body}
|
||||
<footer>
|
||||
<p><strong>How accuracy is attributed.</strong> A round trip is a property of the encoder and
|
||||
decoder together, so to read one stage the other is pinned to that model's 16-bit chunked
|
||||
engine. The encoder chart varies the encoder; the decoder chart varies the decoder.</p>
|
||||
<p><strong>Read these as comparisons, not verdicts.</strong> Compare lines to each other; do
|
||||
not read the absolute height as a quality statement, and do not let any of it substitute for
|
||||
listening.</p>
|
||||
<p><strong>It is not the limiter.</strong> The decoders bake a 0.977 ceiling and real masters
|
||||
here peak at 1.019, so level clipping was the obvious suspect. It is not: re-measured with the
|
||||
input scaled from 0.9 down to 0.225 the SNR moves by 0.03 dB, and an optimal scalar
|
||||
gain-match recovers only 0.27 dB.</p>
|
||||
<p><strong>Why accuracy is plotted with content held fixed.</strong> The first version of
|
||||
these charts measured one excerpt per length — the first L latents of the stream — and the lines
|
||||
bounced by several dB, which looked like a length-dependent defect. It was not. At a fixed L,
|
||||
moving the excerpt swings SNR by <strong>9.4–15.0 dB</strong>, while the entire L axis
|
||||
at a fixed offset moved only <strong>2.95 dB</strong>: a longer L is simply a different,
|
||||
longer piece of music, so one-excerpt-per-L plots the music rather than the model. The charts
|
||||
above therefore hold the audio identical at every point and vary only the processing length —
|
||||
the same thing the TFLite rung table does, and it comes out just as flat.</p>
|
||||
<p><strong>The absolute level.</strong> Round-trip SNR on this real-music set is 4.2–4.9 dB
|
||||
for SAME-L and 3.0–3.9 dB for SAME-S. Measured on our own generated audio the same pipeline
|
||||
scores ~16 dB, and published figures for this family span roughly −23 dB (solo piano) to
|
||||
−4 dB (dense circus). Waveform SNR systematically understates a perceptual autoencoder.</p>
|
||||
<p><strong>SAME-S draw-to-draw noise.</strong> The shipped SAME-S decoder has a
|
||||
<code>RandomNormalLike</code> at its bottleneck, so its round trip carries run-to-run variance
|
||||
that belongs to neither precision nor chunking.</p>
|
||||
</footer>
|
||||
</main>""")
|
||||
print(f" wrote {OUT} ({len(rows)} rows)")
|
||||
@@ -0,0 +1,323 @@
|
||||
"""Vanilla PyTorch FP32 inference for SA3 medium DiT + SAME-L decoder.
|
||||
|
||||
Drop-in alternative to SA3Inference for the "GT PyTorch eager" comparison
|
||||
backend in the gradio app. Uses stable_audio_tools directly for everything
|
||||
except T5 encoding (we still reuse the TRT T5 engine — it's fast, validated
|
||||
equivalent to PT BF16, and avoids re-loading the 538MB T5Gemma weights).
|
||||
|
||||
Public API matches SA3Inference's relevant subset:
|
||||
inf = PTInference() # lazy-loads PT model at first .generate()
|
||||
pcm, timing = inf.generate(prompt, seconds=120.0, steps=8, seed=1)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import json
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
TRT_REPO = Path("/weka2/cj/clod/sa3s/stable-audio-3/optimized/tensorRT")
|
||||
SCRIPTS_DIR = TRT_REPO / "scripts"
|
||||
sys.path.insert(0, str(SCRIPTS_DIR))
|
||||
|
||||
# stable_audio_tools is expected to be importable (the venv this gradio runs in
|
||||
# must include it). The producer venv at /admin/home-cj/sa/venvs/sa3_torch29
|
||||
# has both gradio + tensorrt + stable_audio_tools.
|
||||
SAT_PATH = "/admin/home-cj/sa/stable-audio-tools-dev-latest"
|
||||
if SAT_PATH not in sys.path:
|
||||
sys.path.insert(0, SAT_PATH)
|
||||
|
||||
|
||||
# Defaults — matches the user's "medium / SAME-L / fp32" scope.
|
||||
PT_MODEL_DIR = Path("/weka2/cj/clod/sa3s/models/SA3-M-hf")
|
||||
T5_TRT_PATH = TRT_REPO / "models" / "sm_90" / "t5gemma" / "t5gemma_fp16.trt"
|
||||
SAMEL_CKPT_DIR = Path("/weka2/cj/clod/sa3s/models/SAME-L") # local SAME-L = TRT decoder's weight source
|
||||
SAMES_CKPT_DIR = Path("/weka2/cj/clod/sa3s/models/SAME-S") # local SAME-S = TRT SAME-S decoder's weight source
|
||||
|
||||
SAMPLE_RATE = 44100
|
||||
SAMPLES_PER_LATENT = 4096
|
||||
T5_MAX_LEN = 256
|
||||
SIGMA_MAX = 1.0
|
||||
|
||||
|
||||
def _patch_pt():
|
||||
"""No-op for now. The transformer SDPA fallback in stable_audio_tools is NOT
|
||||
numerically equivalent to flash_attn for this model: encode→decode round-trip
|
||||
cosine drops from 0.78 (flash_attn) → 0.005 (SDPA) on real music. Keep
|
||||
flash_attn enabled."""
|
||||
return
|
||||
|
||||
|
||||
def _resolve_pt_model_dir() -> Path:
|
||||
"""Find the SA3-M HF safetensors dir.
|
||||
|
||||
The producer build references either SA3-M-hf or SA3-medium-ARC. Try both;
|
||||
return the first match.
|
||||
"""
|
||||
for name in ("SA3-M-hf", "SA3-M", "SA3-medium-ARC", "SA3-medium"):
|
||||
d = Path("/weka2/cj/clod/sa3s/models") / name
|
||||
if (d / "model.safetensors").exists() and (d / "model_config.json").exists():
|
||||
return d
|
||||
raise FileNotFoundError(
|
||||
f"Couldn't locate the medium DiT PyTorch checkpoint under "
|
||||
f"/weka2/cj/clod/sa3s/models/. Looked for SA3-M-hf / SA3-M / SA3-medium-ARC / SA3-medium."
|
||||
)
|
||||
|
||||
|
||||
def _load_pt_model(model_dir: Path, device="cuda"):
|
||||
from stable_audio_tools.models.factory import create_model_from_config
|
||||
from stable_audio_tools.models.utils import copy_state_dict
|
||||
from safetensors.torch import load_file
|
||||
|
||||
with open(model_dir / "model_config.json") as f:
|
||||
cfg = json.load(f)
|
||||
# SAT's T5GemmaConditioner only takes model_name; HF cfg adds repo_id/subfolder
|
||||
for c in cfg.get("model", {}).get("conditioning", {}).get("configs", []):
|
||||
if c.get("type") == "t5gemma":
|
||||
cc = c.get("config", {})
|
||||
cc.pop("repo_id", None); cc.pop("subfolder", None)
|
||||
cc.setdefault("model_name", "google/t5gemma-b-b-ul2")
|
||||
model = create_model_from_config(cfg)
|
||||
sd = load_file(model_dir / "model.safetensors")
|
||||
copy_state_dict(model, sd)
|
||||
model = model.float().to(device).eval()
|
||||
return model, sd
|
||||
|
||||
|
||||
def _load_decoder(ckpt_dir: Path, name: str, device="cuda"):
|
||||
"""Load a LOCAL SAME-{S,L} AudioAutoencoder — same weights as the TRT decoder.
|
||||
|
||||
Noise sources are disabled to match TRT's deterministic rewrite (the producer's
|
||||
ONNX validator disables `bottleneck.noise_regularize` and
|
||||
`decoder.layers[3].mask_noise` before comparing — those random kernels would
|
||||
inject buzz at inference if left on).
|
||||
"""
|
||||
from stable_audio_tools.models.factory import create_model_from_config
|
||||
from stable_audio_tools.models.utils import load_ckpt_state_dict, copy_state_dict
|
||||
with open(ckpt_dir / f"{name}.json") as f:
|
||||
cfg = json.load(f)
|
||||
model = create_model_from_config(cfg).to(device).eval()
|
||||
sd = load_ckpt_state_dict(str(ckpt_dir / f"{name}.ckpt"))
|
||||
copy_state_dict(model, sd)
|
||||
model.bottleneck.noise_regularize = False
|
||||
model.decoder.layers[3].mask_noise = 0
|
||||
return model
|
||||
|
||||
|
||||
def _load_samel_decoder(device="cuda"):
|
||||
return _load_decoder(SAMEL_CKPT_DIR, "SAME-L", device=device)
|
||||
|
||||
|
||||
def _load_sames_decoder(device="cuda"):
|
||||
return _load_decoder(SAMES_CKPT_DIR, "SAME-S", device=device)
|
||||
|
||||
|
||||
def _build_dist_shift():
|
||||
"""Import runtime.DistributionShift for the sigma schedule shift."""
|
||||
import runtime as rt
|
||||
return rt.DistributionShift()
|
||||
|
||||
|
||||
class PTInference:
|
||||
"""Vanilla PyTorch FP32 inference: PT DiT + PT decoder(s), TRT T5.
|
||||
|
||||
Scope: medium DiT + SAME-L / SAME-S decoders + FP32 only.
|
||||
The DiT is loaded once; both decoders are loaded so generate(decoder=...)
|
||||
can pick between them without reloading anything.
|
||||
Other configs raise NotImplementedError.
|
||||
|
||||
Thread-safe via an internal Lock (the underlying PyTorch model state is
|
||||
shared mutable; serialize generate() calls).
|
||||
"""
|
||||
|
||||
def __init__(self, device: str = "cuda", *,
|
||||
load_samel: bool = True, load_sames: bool = True):
|
||||
import threading
|
||||
self._lock = threading.Lock()
|
||||
self.device = device
|
||||
|
||||
print(" PT FP32 GT: patching stable_audio_tools + loading checkpoint...", flush=True)
|
||||
_patch_pt()
|
||||
t0 = time.time()
|
||||
model_dir = _resolve_pt_model_dir()
|
||||
self.pt_model, self.sd = _load_pt_model(model_dir, device=device)
|
||||
print(f" PT model loaded from {model_dir.name} in {time.time()-t0:.1f}s "
|
||||
f"({sum(p.numel() for p in self.pt_model.parameters())/1e6:.0f}M params)",
|
||||
flush=True)
|
||||
self.pt_dit = self.pt_model.model.model
|
||||
|
||||
# Use the LOCAL SAME-L / SAME-S AudioAutoencoders — same weights as the
|
||||
# TRT decoder engines (the engines are an ONNX rewrite of these PT models).
|
||||
self.pt_samel = None
|
||||
self.pt_sames = None
|
||||
if load_samel:
|
||||
t0 = time.time()
|
||||
self.pt_samel = _load_samel_decoder(device=device)
|
||||
print(f" SAME-L decoder loaded from {SAMEL_CKPT_DIR.name} in {time.time()-t0:.1f}s "
|
||||
f"(noise sources disabled)", flush=True)
|
||||
if load_sames:
|
||||
t0 = time.time()
|
||||
self.pt_sames = _load_sames_decoder(device=device)
|
||||
print(f" SAME-S decoder loaded from {SAMES_CKPT_DIR.name} in {time.time()-t0:.1f}s "
|
||||
f"(noise sources disabled)", flush=True)
|
||||
|
||||
# Reuse TRT T5 (fast + validated). Loaded lazily via canon to share
|
||||
# the same heavy-imports path the SA3Inference uses.
|
||||
import sa3_trt_core as canon
|
||||
canon._import_heavy()
|
||||
import runtime as rt
|
||||
rt.MODELS_DIR = str(TRT_REPO / "models")
|
||||
rt.ARCH_DIR = str(TRT_REPO / "models" / "sm_90")
|
||||
self.tokenizer = rt.load()["tokenizer"]
|
||||
self.t5_runner = canon.TRTRunner(T5_TRT_PATH)
|
||||
self.dist_shift = _build_dist_shift()
|
||||
self._canon = canon
|
||||
|
||||
# ── conditioning helper (FP32 PT path with the 2π factor baked in) ────
|
||||
def _build_cond(self, t5_hidden, t5_mask, seconds_total: float):
|
||||
sd = self.sd
|
||||
padding_emb = sd["conditioner.conditioners.prompt.padding_embedding"].float().to(self.device)
|
||||
sec_w = sd["conditioner.conditioners.seconds_total.embedder.embedding.1.weight"].float().to(self.device)
|
||||
sec_b = sd["conditioner.conditioners.seconds_total.embedder.embedding.1.bias"].float().to(self.device)
|
||||
half = 128
|
||||
ramp = torch.linspace(0, 1, half, device=self.device)
|
||||
freqs = torch.exp(ramp * (math.log(10000.0) - math.log(0.5)) + math.log(0.5))
|
||||
s = torch.tensor([seconds_total], device=self.device).clamp(0, 384) / 384
|
||||
args = s.unsqueeze(-1) * freqs * 2 * math.pi
|
||||
ff = torch.cat([args.cos(), args.sin()], dim=-1)
|
||||
sec_emb = (ff @ sec_w.T + sec_b).unsqueeze(1)
|
||||
m = t5_mask.unsqueeze(-1).bool()
|
||||
pe = padding_emb.view(1, 1, -1).expand_as(t5_hidden)
|
||||
t5_padded = torch.where(m, t5_hidden, pe)
|
||||
return torch.cat([t5_padded, sec_emb], dim=1), sec_emb.squeeze(1)
|
||||
|
||||
@staticmethod
|
||||
def resolve_T_lat(seconds: float, decoder: str) -> int:
|
||||
"""seconds → T_lat (no even-bump for SAME-L, since SAME-L doesn't care)."""
|
||||
T_lat = max(1, math.ceil(seconds * SAMPLE_RATE / SAMPLES_PER_LATENT))
|
||||
return T_lat
|
||||
|
||||
def generate(self, prompt: str, *,
|
||||
seconds: float = 120.0, steps: int = 8,
|
||||
seed: Optional[int] = None,
|
||||
decoder: str = "same-l",
|
||||
# Future kwargs (NotImplementedError until wired):
|
||||
init_noise_level: float = 1.0,
|
||||
negative_prompt: Optional[str] = None,
|
||||
cfg: float = 1.0,
|
||||
init_audio_path: Optional[str] = None,
|
||||
inpaint_range: Optional[tuple] = None,
|
||||
) -> tuple[np.ndarray, dict]:
|
||||
if cfg != 1.0:
|
||||
raise NotImplementedError("CFG not yet wired through PTInference")
|
||||
if init_audio_path is not None or inpaint_range is not None:
|
||||
raise NotImplementedError("audio-to-audio / inpaint not yet wired")
|
||||
if init_noise_level != 1.0:
|
||||
raise NotImplementedError("non-unity init_noise_level not yet wired")
|
||||
if decoder == "same-l":
|
||||
pt_decoder = self.pt_samel
|
||||
elif decoder == "same-s":
|
||||
pt_decoder = self.pt_sames
|
||||
else:
|
||||
raise ValueError(f"unknown decoder={decoder!r}; valid: same-l, same-s")
|
||||
if pt_decoder is None:
|
||||
raise RuntimeError(f"decoder {decoder!r} was not loaded — pass load_{decoder.replace('-','_')}=True to PTInference()")
|
||||
|
||||
canon = self._canon
|
||||
device = self.device
|
||||
T_lat = self.resolve_T_lat(seconds, decoder)
|
||||
if not (1 <= T_lat <= 4096):
|
||||
raise ValueError(f"T_lat={T_lat} out of [1, 4096]")
|
||||
if seed is None:
|
||||
import random
|
||||
seed = random.randint(0, 2**31 - 1)
|
||||
|
||||
with self._lock:
|
||||
t_total = time.time()
|
||||
|
||||
# Match official generate_diffusion_cond: disable TF32 + benchmark for
|
||||
# FP32-faithful math.
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
# 1. T5 encode via TRT (fast).
|
||||
t0 = time.time()
|
||||
t5_hidden, t5_mask = canon.t5gemma_encode(self.t5_runner, self.tokenizer, prompt)
|
||||
t5_hidden = t5_hidden.to(device).float(); t5_mask = t5_mask.to(device)
|
||||
t5_ms = (time.time() - t0) * 1000
|
||||
|
||||
# 2. Conditioning + sigmas.
|
||||
cond_ca, cond_g = self._build_cond(t5_hidden, t5_mask, seconds)
|
||||
sigmas = torch.linspace(SIGMA_MAX, 0.0, steps + 1, device=device)
|
||||
sigmas = self.dist_shift.shift(sigmas, T_lat); sigmas[0] = SIGMA_MAX
|
||||
|
||||
# 3. Sampling.
|
||||
g_init = torch.Generator(device=device); g_init.manual_seed(int(seed))
|
||||
g_loop = torch.Generator(device=device); g_loop.manual_seed(int(seed) + 1)
|
||||
x = torch.randn(1, 256, T_lat, device=device, dtype=torch.float32, generator=g_init)
|
||||
local_add_cond = torch.zeros(1, 257, T_lat, device=device)
|
||||
|
||||
t0 = time.time()
|
||||
with torch.no_grad():
|
||||
for i in range(steps):
|
||||
t_curr = sigmas[i]; t_next = sigmas[i + 1]
|
||||
v = self.pt_dit._forward(
|
||||
x, t_curr.unsqueeze(0).contiguous(),
|
||||
cross_attn_cond=cond_ca, global_embed=cond_g,
|
||||
local_add_cond=local_add_cond,
|
||||
)
|
||||
denoised = x - t_curr * v
|
||||
if i < steps - 1:
|
||||
noise = torch.randn(*x.shape, device=device, dtype=x.dtype, generator=g_loop)
|
||||
x = (1.0 - t_next) * denoised + t_next * noise
|
||||
else:
|
||||
x = denoised
|
||||
torch.cuda.synchronize()
|
||||
sampling_ms = (time.time() - t0) * 1000
|
||||
|
||||
# 4. Decode FP32 via selected decoder (matches TRT decoder weights).
|
||||
t0 = time.time()
|
||||
with torch.no_grad():
|
||||
audio_fp32 = pt_decoder.decode(x.float())
|
||||
torch.cuda.synchronize()
|
||||
decode_ms = (time.time() - t0) * 1000
|
||||
|
||||
# 5. FP32 audio → int16 stereo PCM (same formula as TRT pipeline).
|
||||
pcm_torch = (audio_fp32.clamp(-1.0, 1.0) * 32767.0).to(torch.int16)
|
||||
pcm = pcm_torch.squeeze(0).T.contiguous().cpu().numpy()
|
||||
|
||||
# Trim to requested seconds.
|
||||
actual_samples = int(round(seconds * SAMPLE_RATE))
|
||||
pcm = pcm[:actual_samples].copy()
|
||||
|
||||
inference_ms = (time.time() - t_total) * 1000
|
||||
return pcm, {
|
||||
"inference_ms": inference_ms,
|
||||
"graph_build_ms": 0.0, # n/a for PT
|
||||
"realtime": (seconds * 1000.0) / inference_ms if inference_ms > 0 else 0.0,
|
||||
"seed": int(seed),
|
||||
"T_lat": T_lat,
|
||||
"samples": actual_samples,
|
||||
"t5_ms": t5_ms,
|
||||
"sampling_ms": sampling_ms,
|
||||
"decode_ms": decode_ms,
|
||||
}
|
||||
|
||||
|
||||
# ── module-level singleton (lazy) ─────────────────────────────────────────
|
||||
_pt_inference: Optional[PTInference] = None
|
||||
|
||||
|
||||
def get_pt_inference() -> PTInference:
|
||||
"""Return the singleton PTInference, loading on first call (~30 sec)."""
|
||||
global _pt_inference
|
||||
if _pt_inference is None:
|
||||
_pt_inference = PTInference()
|
||||
return _pt_inference
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
Runtime support: tokenizer + distribution-shift constructors.
|
||||
|
||||
Everything model-dependent (tokenizer, T5Gemma engine, DiT engines, decoder
|
||||
engines) ships from HuggingFace under models/. The DiT engines bundle their
|
||||
per-DiT conditioner tensors (padding_embedding + seconds_total Linear) as
|
||||
graph Constants, and DistributionShift's schedule constants are SA3-canonical
|
||||
and live as defaults on the class.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
|
||||
# runtime.py lives in scripts/; models/ lives one level up at the repo root.
|
||||
_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
MODELS_DIR = os.path.join(_REPO, "models")
|
||||
# ARCH_DIR is set by sa3_trt.py before calling load() (it knows the GPU arch).
|
||||
ARCH_DIR = MODELS_DIR # placeholder; overwritten by sa3_trt.py
|
||||
IO_CHANNELS = 256
|
||||
SAMPLE_RATE = 44100
|
||||
DOWNSAMPLING_RATIO = 4096
|
||||
|
||||
|
||||
class FastTokenizer:
|
||||
"""Tokenizer using HuggingFace tokenizers lib directly (no transformers import)."""
|
||||
def __init__(self, tokenizer_json_path, max_length=256, pad_token_id=0):
|
||||
from tokenizers import Tokenizer
|
||||
self.tokenizer = Tokenizer.from_file(tokenizer_json_path)
|
||||
self.max_length = max_length
|
||||
self.pad_token_id = pad_token_id
|
||||
|
||||
def __call__(self, text, return_tensors="pt", max_length=None, padding="max_length", truncation=True):
|
||||
max_len = max_length or self.max_length
|
||||
enc = self.tokenizer.encode(text)
|
||||
ids = enc.ids[:max_len]
|
||||
mask = [1] * len(ids)
|
||||
if padding == "max_length":
|
||||
pad_len = max_len - len(ids)
|
||||
ids = ids + [self.pad_token_id] * pad_len
|
||||
mask = mask + [0] * pad_len
|
||||
if return_tensors == "pt":
|
||||
return {
|
||||
"input_ids": torch.tensor([ids], dtype=torch.long),
|
||||
"attention_mask": torch.tensor([mask], dtype=torch.long),
|
||||
}
|
||||
return {"input_ids": ids, "attention_mask": mask}
|
||||
|
||||
|
||||
def load_tokenizer():
|
||||
"""Load tokenizer using fast tokenizers lib (no HF transformers import).
|
||||
|
||||
tokenizer.json is bundled in the repo next to this module (it's arch-
|
||||
agnostic — same T5Gemma vocab whether you're on sm_90, sm_100, sm_120
|
||||
etc.). Falls back to the legacy per-arch location if the bundled file
|
||||
is somehow missing (e.g. older checkouts).
|
||||
"""
|
||||
bundled = os.path.join(os.path.dirname(__file__), "tokenizer.json")
|
||||
legacy = os.path.join(ARCH_DIR, "t5gemma", "tokenizer.json")
|
||||
path = bundled if os.path.exists(bundled) else legacy
|
||||
return FastTokenizer(path)
|
||||
|
||||
|
||||
class DistributionShift:
|
||||
"""FluxDistributionShift — schedule warp used by SA3 models (rf_denoiser).
|
||||
|
||||
Without this warp the model is conditioned on the wrong t at each step and
|
||||
final latent magnitudes blow up ~60×. Defaults match HF SA3 model configs
|
||||
(all variants share base_shift=0.5, max_shift=1.15, min/max_length=256/4096).
|
||||
"""
|
||||
def __init__(self, base_shift=0.5, max_shift=1.15, min_length=256, max_length=4096):
|
||||
self.base_shift = base_shift
|
||||
self.max_shift = max_shift
|
||||
self.min_length = min_length
|
||||
self.max_length = max_length
|
||||
|
||||
def shift(self, t, seq_len):
|
||||
sl = min(max(int(seq_len), self.min_length), self.max_length)
|
||||
mu = -(self.base_shift + (self.max_shift - self.base_shift) *
|
||||
(sl - self.min_length) / (self.max_length - self.min_length))
|
||||
return 1 - math.exp(mu) / (math.exp(mu) + (1.0 / (1.0 - t) - 1.0))
|
||||
|
||||
|
||||
def load():
|
||||
"""Load the runtime pieces needed for TRT inference: tokenizer and distribution
|
||||
shift. Neither requires the full PyTorch stable_audio_tools model.
|
||||
"""
|
||||
return {"tokenizer": load_tokenizer(), "dist_shift": DistributionShift()}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
state = load()
|
||||
print(f"Loaded: {list(state.keys())}")
|
||||
@@ -0,0 +1,724 @@
|
||||
"""SA3 TRT — gradio web UI (debug build).
|
||||
|
||||
Adds (vs the MVP):
|
||||
- Model picker: sm-music / sm-sfx / medium (hot-swap with cached SA3Inference
|
||||
instances; first switch ~5-7s, subsequent instant)
|
||||
- Engine-variant picker per model: auto-detects all `dit_*.trt` files in the
|
||||
model dir (canonical fixed, archived buggy, fp32, etc.) so you can A/B
|
||||
different precisions / quantizations without restarting
|
||||
- Spectrogram display: Underfit-style 3-band tinted stereo mel spectrogram
|
||||
rendered inline alongside the audio
|
||||
|
||||
Launch:
|
||||
./sa3-gradio # share=True by default, sm-music + same-s
|
||||
./sa3-gradio --dit medium
|
||||
./sa3-gradio --no-share # local-only
|
||||
|
||||
The previously-required `--dit ...` is now just the *initial* model — the
|
||||
runtime dropdown lets you switch between variants/models without restart.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse
|
||||
import base64
|
||||
import math
|
||||
import sys
|
||||
import html as html_lib
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
SCRIPTS_DIR = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(SCRIPTS_DIR))
|
||||
|
||||
from sa3_trt import SA3Inference, SAMPLE_RATE, SAMPLES_PER_LATENT # noqa: E402
|
||||
import sa3_trt_core as canon # noqa: E402
|
||||
from spec import render_spectrogram_png # noqa: E402
|
||||
|
||||
|
||||
OUTPUT_DIR = SCRIPTS_DIR.parent / "output" / "gradio"
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
KEEP_RECENT_N = 20
|
||||
|
||||
from gradio_ui import ( # backend-agnostic UI layer, forked from optimized/mlx
|
||||
# The other _JS_* constants (promote / pause-others / scroll / playhead) are wired
|
||||
# inside render_player and render_history themselves, so they stay private to that
|
||||
# module; only the two this file attaches to its own events are imported.
|
||||
_JS_FIX_SLIDERS, _save_wav, render_history, render_player, render_queue_status, verbose_basename,
|
||||
)
|
||||
|
||||
MODELS_ROOT = SCRIPTS_DIR.parent / "models" / canon.ARCH
|
||||
|
||||
|
||||
# Default decoder for each DiT bundle.
|
||||
DEFAULT_DECODERS = {"sm-music": "same-s", "sm-sfx": "same-s", "medium": "same-l"}
|
||||
|
||||
# DiT-name → on-disk subdir mapping (matches sa3_trt_core's DIT_CHOICES)
|
||||
DIT_SUBDIRS = {"sm-music": "sa3-sm-music", "sm-sfx": "sa3-sm-sfx", "medium": "sa3-m"}
|
||||
|
||||
|
||||
# Magic marker (matched against str(variant_path)) telling generate() to
|
||||
# dispatch to the PyTorch-eager backend instead of TRT. Only offered for
|
||||
# the medium+SAME-L+fp32 combo per the user's scope.
|
||||
PT_EAGER_VARIANT = "<PT_EAGER_FP32_GT>"
|
||||
|
||||
|
||||
def discover_variants(dit_name: str) -> list[tuple[str, Path]]:
|
||||
"""Return [(label, path)] of available DiT engine files for this model.
|
||||
|
||||
Scans models/<arch>/<dit_subdir>/dit_*.trt. The canonical engine
|
||||
(dit_fp16.trt) is always first if present; other variants follow
|
||||
alphabetically. For the medium DiT also appends a pseudo-variant
|
||||
"pytorch fp32 (GT)" that dispatches to the PyTorch-eager backend.
|
||||
"""
|
||||
subdir = DIT_SUBDIRS.get(dit_name)
|
||||
if subdir is None:
|
||||
return []
|
||||
d = MODELS_ROOT / subdir
|
||||
if not d.exists():
|
||||
return []
|
||||
files = [f for f in sorted(d.glob("dit_*.trt")) if f.name not in RETIRED_DITS]
|
||||
canonical_name = "dit_fp16.trt"
|
||||
canonical = [f for f in files if f.name == canonical_name]
|
||||
others = [f for f in files if f.name != canonical_name]
|
||||
out = []
|
||||
for f in canonical + others:
|
||||
label = f.name[len("dit_"):-len(".trt")] if f.name.startswith("dit_") and f.name.endswith(".trt") else f.name
|
||||
if f.name == canonical_name:
|
||||
label = "fp16 (canonical)"
|
||||
elif "buggy" in f.name:
|
||||
label = label + " ← old, broken"
|
||||
out.append((label, f))
|
||||
# Pseudo-variant: PT FP32 GT (currently medium only).
|
||||
if dit_name == "medium":
|
||||
out.append(("pytorch fp32 GT (slow, vanilla eager)", Path(PT_EAGER_VARIANT)))
|
||||
return out
|
||||
|
||||
|
||||
# Pre-chunkable engines. Removed from the model repo's download path and hidden from the picker.
|
||||
RETIRED_DECODERS = {"dec_dynamic_bf16.trt", "dec_dynamic_triton_swa.trt", "dec_dynamic_fp32.trt",
|
||||
"dec_fp8.trt", "dec_fp8_fast.trt", "dec_dynamic_fp16mixed.trt"}
|
||||
# Experimental DiT engines that are not part of the shipped set. "fp16mixed" is the retired name
|
||||
# for plain fp16 -- every tier is mixed-precision, so the qualifier only ever meant "old".
|
||||
RETIRED_DITS = {"dit_fp16mixed.trt", "dit_fp16_bands4096_32768.trt", "dit_fp16_maxL32768.trt"}
|
||||
|
||||
|
||||
def _precision_of(filename: str) -> str:
|
||||
"""Pull the precision token out of an engine filename: dec_fp16_chunkable_limiter -> fp16."""
|
||||
for tok in ("fp32", "fp16", "bf16", "fp8"):
|
||||
if f"_{tok}" in filename:
|
||||
return tok
|
||||
return ""
|
||||
|
||||
|
||||
def discover_decoder_variants(decoder_name: str) -> list[tuple[str, Path]]:
|
||||
"""Return [(label, path)] of decoder quantization tiers for this decoder —
|
||||
canonical first. Uses the known tier set (canonical / fp8; see quantize/README.md); engines
|
||||
auto-download from HF on selection if missing. Any extra local dec_*.trt not in the known set
|
||||
are appended, so a self-built variant shows up — except the RETIRED pre-chunkable engines,
|
||||
which are filtered out. A leftover copy of one of those in models/ would otherwise appear in
|
||||
the picker and silently serve pre-limiter audio with a 5.7-8.1 GB scratch reservation.
|
||||
"""
|
||||
out, known = [], set(RETIRED_DECODERS)
|
||||
for tier, fname in canon.DECODER_TIER_FILENAME.get(decoder_name, {}).items():
|
||||
# Name the tier by the precision it actually is -- SAME-L canonical is fp16, SAME-S is
|
||||
# bf16 -- rather than rendering "canonical (canonical)".
|
||||
prec = _precision_of(fname)
|
||||
label = f"{prec} (canonical)" if tier == "canonical" else (prec or tier)
|
||||
out.append((label, canon.ARCH_DIR / decoder_name / fname))
|
||||
known.add(fname)
|
||||
d = canon.ARCH_DIR / decoder_name
|
||||
if d.exists():
|
||||
for f in sorted(d.glob("dec_*.trt")):
|
||||
if f.name in known:
|
||||
continue
|
||||
stem = f.name
|
||||
label = (stem[len("dec_dynamic_"):-len(".trt")] if stem.startswith("dec_dynamic_")
|
||||
else stem[len("dec_"):-len(".trt")])
|
||||
out.append((label, f))
|
||||
return out
|
||||
|
||||
|
||||
def _prune_old_outputs():
|
||||
wavs = sorted(OUTPUT_DIR.glob("*.wav"), key=lambda p: p.stat().st_mtime, reverse=True)
|
||||
for old in wavs[KEEP_RECENT_N:]:
|
||||
try:
|
||||
old.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
# ── SA3Inference cache (model+variant → instance) ──────────────────────────
|
||||
# Each cached entry holds 5-15 GB of VRAM (engines + persistent graph buffers).
|
||||
# H100 80GB fits ~5 sm-music or ~3 medium combos before OOM. We LRU-evict to
|
||||
# stay under that limit. Switching back to an evicted variant re-loads (~5-7s).
|
||||
_inference_cache: dict[tuple[str, str, str, str, bool], SA3Inference] = {}
|
||||
_inference_lru: list[tuple[str, str, str, str]] = [] # MRU at end
|
||||
_INFERENCE_CACHE_MAX = 2
|
||||
|
||||
|
||||
def _mb(b):
|
||||
return f"{b / 2**20:,.0f} MB" if b else "—"
|
||||
|
||||
|
||||
def render_debug(inf, t_fused, t_stage):
|
||||
"""Per-stage ms + VRAM, and what the fused mega-graph costs by comparison.
|
||||
|
||||
⚠ The stage split is measured on the UNFUSED path, because a captured CUDA graph
|
||||
cannot be instrumented from inside: torch rejects Event.record() during capture, and
|
||||
even cudaEventRecordWithFlags(..., cudaEventRecordExternal) is accepted only to fail
|
||||
at capture_end(). Adding timers would mean splitting the mega-graph into per-stage
|
||||
graphs, which would change the very thing being measured. So debug mode times the
|
||||
same engines run stage-by-stage and reports the fused total beside it — the gap
|
||||
between the two IS the fusion win.
|
||||
|
||||
VRAM is split the same way it is actually paid: engine weights and scratch are
|
||||
resident for the whole process, while `transient` is what a stage allocates on top.
|
||||
"""
|
||||
fp = inf.engine_footprint()
|
||||
sv = t_stage.get("stage_vram", {})
|
||||
rows = [("encoder", t_stage.get("encode_ms", 0.0), sv.get("encode"), fp.get("enc")),
|
||||
("T5Gemma", t_stage.get("t5_ms", 0.0), sv.get("t5"), fp.get("t5")),
|
||||
(f"DiT x{t_stage.get('steps', 0)}",
|
||||
t_stage.get("sampling_ms", 0.0), sv.get("dit"), fp.get("dit")),
|
||||
("decoder", t_stage.get("decode_ms", 0.0), sv.get("decode"), fp.get("dec"))]
|
||||
th = ("padding:3px 10px 3px 0; text-align:right; font-weight:600; "
|
||||
"border-bottom:1px solid rgba(127,127,127,0.3)")
|
||||
td = "padding:2px 10px 2px 0; text-align:right; font-variant-numeric:tabular-nums"
|
||||
out = ['<div style="font-size:0.82em; margin-top:6px">',
|
||||
'<table style="border-collapse:collapse">',
|
||||
f'<tr><th style="{th}; text-align:left">stage</th><th style="{th}">ms</th>'
|
||||
f'<th style="{th}">transient</th><th style="{th}">weights</th>'
|
||||
f'<th style="{th}">scratch</th></tr>']
|
||||
for name, ms, tv, f in rows:
|
||||
if f is None and not ms:
|
||||
continue # encoder absent unless a2a/inpaint ran
|
||||
sc = _mb(f["scratch"]) + (" (shared)" if f and f["shared"] else "") if f else "—"
|
||||
out.append(f'<tr><td style="{td}; text-align:left">{html_lib.escape(name)}</td>'
|
||||
f'<td style="{td}">{ms:,.1f}</td><td style="{td}">{_mb(tv)}</td>'
|
||||
f'<td style="{td}">{_mb(f["weights"]) if f else "—"}</td>'
|
||||
f'<td style="{td}">{sc}</td></tr>')
|
||||
# Whatever the four stages do not account for: schedule build, noise, the int16 cast
|
||||
# and the GPU->CPU copy. Shown rather than hidden, so the column sums to the total.
|
||||
unfused = t_stage.get("inference_ms", 0.0)
|
||||
other = unfused - sum(r[1] for r in rows)
|
||||
if other > 0.5:
|
||||
out.append(f'<tr><td style="{td}; text-align:left; opacity:0.7">other '
|
||||
f'<span style="opacity:0.7">(schedule, noise, int16 cast, GPU→CPU)</span>'
|
||||
f'</td><td style="{td}; opacity:0.7">{other:,.1f}</td>'
|
||||
f'<td style="{td}">—</td><td style="{td}">—</td>'
|
||||
f'<td style="{td}">—</td></tr>')
|
||||
out.append("</table>")
|
||||
shared = fp.get("_shared_scratch", 0)
|
||||
note = []
|
||||
if t_fused.get("eager"):
|
||||
note.append("this render used the eager path (CFG / audio-to-audio / inpaint), "
|
||||
"so the breakdown IS the render — no second pass was run")
|
||||
else:
|
||||
fused = t_fused.get("inference_ms", 0.0)
|
||||
gain = f"{unfused / fused:.2f}x faster" if fused > 0 else ""
|
||||
note.append(f"unfused {unfused:,.1f} ms · "
|
||||
f"<b>fused mega-graph {fused:,.1f} ms</b> ({gain})")
|
||||
note.append(f"shared scratch {_mb(shared)} — one buffer for every engine above, so the "
|
||||
f"scratch column does not sum")
|
||||
note.append("a captured graph allocates nothing at replay, so the mega-graph's peak "
|
||||
"VRAM equals its resident footprint")
|
||||
out += [f'<div style="color:#888; margin-top:4px">{n}</div>' for n in note]
|
||||
out.append("</div>")
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _evict_inference(key):
|
||||
"""Free the engines + caches held by a cached SA3Inference before dropping it.
|
||||
|
||||
The teardown ORDER is load-bearing and lives in SA3Inference.close(): captured CUDA
|
||||
graphs and the eager contexts must go before the runners whose contexts and shared
|
||||
scratch they point into. This used to free the runners first and clear the graphs
|
||||
after, which poisoned the CUDA context — the next render in a *different* instance
|
||||
died with `an illegal memory access was encountered`, so switching models past the
|
||||
cache limit bricked the server until restart.
|
||||
"""
|
||||
inf = _inference_cache.pop(key, None)
|
||||
if inf is None:
|
||||
return
|
||||
try:
|
||||
inf.close()
|
||||
except Exception as e:
|
||||
print(f" (eviction cleanup raised {type(e).__name__}: {e})")
|
||||
del inf
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
|
||||
# Set once from --chunking in main(). A module-level default keeps it out of every call site
|
||||
# and out of the UI, which is right: it is a deployment choice (memory vs speed), not a per-render
|
||||
# one -- switching it means selecting a different optimization profile, so the engines reload.
|
||||
_CHUNKING = True
|
||||
|
||||
|
||||
def get_inference(dit: str, decoder: str, dit_variant_path: str,
|
||||
dec_variant_path: str,
|
||||
default_T_lat: int, default_steps: int,
|
||||
default_seconds: float, quiet: bool,
|
||||
chunking: bool | None = None) -> SA3Inference:
|
||||
"""Return a warm SA3Inference, building (and caching) if not yet loaded.
|
||||
|
||||
The cache key is (dit, decoder, dit_variant_path, dec_variant_path) so
|
||||
swapping either the DiT *or* the decoder variant triggers a fresh load.
|
||||
LRU eviction caps the cache at _INFERENCE_CACHE_MAX entries.
|
||||
"""
|
||||
chunking = _CHUNKING if chunking is None else chunking
|
||||
# chunking is part of the key: it selects the profile, so the two modes are different
|
||||
# contexts with different scratch and cannot share a cached SA3Inference.
|
||||
key = (dit, decoder, dit_variant_path, dec_variant_path, chunking)
|
||||
if key in _inference_cache:
|
||||
if key in _inference_lru:
|
||||
_inference_lru.remove(key)
|
||||
_inference_lru.append(key)
|
||||
return _inference_cache[key]
|
||||
|
||||
# LRU evict BEFORE loading new — frees VRAM so the new load doesn't OOM.
|
||||
while len(_inference_cache) >= _INFERENCE_CACHE_MAX:
|
||||
oldest = _inference_lru.pop(0)
|
||||
print(f"\n ← LRU-evicting SA3Inference{oldest[:2]} "
|
||||
f"dit={Path(oldest[2]).name} dec={Path(oldest[3]).name}")
|
||||
_evict_inference(oldest)
|
||||
|
||||
# Override the canonical engine path lookups before constructing.
|
||||
canon.DIT_CHOICES[dit]["engine"] = Path(dit_variant_path)
|
||||
canon.DECODER_PATHS[decoder] = Path(dec_variant_path)
|
||||
dec_tier = canon.decoder_tier_from_filename(decoder, Path(dec_variant_path).name)
|
||||
print(f"\n → loading SA3Inference({dit!r}, {decoder!r}, "
|
||||
f"dit={Path(dit_variant_path).name}, "
|
||||
f"dec={Path(dec_variant_path).name}, tier={dec_tier})")
|
||||
inf = SA3Inference(dit, decoder,
|
||||
dec_precision=dec_tier,
|
||||
chunking=chunking,
|
||||
# pass the picked files explicitly -- mutating canon.DIT_CHOICES /
|
||||
# DECODER_PATHS above does not reach the resolver, so the dropdowns
|
||||
# were silently ignored for the DiT.
|
||||
dit_engine=(None if str(dit_variant_path).startswith("<")
|
||||
else dit_variant_path),
|
||||
dec_engine=dec_variant_path,
|
||||
default_T_lat=default_T_lat,
|
||||
default_steps=default_steps,
|
||||
default_seconds=default_seconds,
|
||||
quiet=quiet)
|
||||
_inference_cache[key] = inf
|
||||
_inference_lru.append(key)
|
||||
return inf
|
||||
|
||||
|
||||
# ── Gradio UI ──────────────────────────────────────────────────────────────
|
||||
_css = ("#sa3-promote{display:none !important}"
|
||||
"#sa3-out{gap:4px !important}"
|
||||
"#sa3-out .html-container{padding:0 !important; margin:0 !important}"
|
||||
# the queued chip is clickable once a take is ready — say so on hover/focus
|
||||
".sa3-queued{transition:background 0.12s, color 0.12s}"
|
||||
".sa3-queued:hover,.sa3-queued:focus{background:rgba(127,127,127,0.22) !important;"
|
||||
"color:#bbb !important; outline:none}")
|
||||
|
||||
|
||||
def build_ui(initial_dit: str, initial_decoder: str, *,
|
||||
share: bool, quiet: bool,
|
||||
default_T_lat: int, default_steps: int, default_seconds: float):
|
||||
"""Same UI as the MLX gradio -- history, a pre-rendered next take, radio mode, the rich
|
||||
player -- over the TensorRT backend, plus the controls only TRT has (engine variant per
|
||||
model, decoder quantisation tier, chunked vs single-shot).
|
||||
|
||||
No LoRA panel: the TRT engines carry no LoRA branches, so there is nothing to hot-swap
|
||||
without first building refit-capable or branch-capable engines.
|
||||
"""
|
||||
import gradio as gr
|
||||
|
||||
initial_variants = discover_variants(initial_dit)
|
||||
if not initial_variants:
|
||||
raise RuntimeError(f"no DiT engines found for {initial_dit} under {MODELS_ROOT}")
|
||||
initial_variant_path = str(initial_variants[0][1])
|
||||
initial_dec_variants = discover_decoder_variants(initial_decoder)
|
||||
if not initial_dec_variants:
|
||||
raise RuntimeError(f"no decoder engines found for {initial_decoder} under {MODELS_ROOT}")
|
||||
initial_dec_variant_path = str(initial_dec_variants[0][1])
|
||||
get_inference(initial_dit, initial_decoder, initial_variant_path,
|
||||
initial_dec_variant_path, default_T_lat, default_steps,
|
||||
default_seconds, quiet)
|
||||
|
||||
DIT_OPTIONS = list(DEFAULT_DECODERS.keys())
|
||||
# The DiT profile runs L=1..4096 and every shipped chunkable autoencoder reports a
|
||||
# floor of 1 latent, so 1 s is the honest slider minimum. It used to be 3 s, derived
|
||||
# from the stale DECODER_MIN_L=32 that SA3Inference._dec_min_L() was returning because
|
||||
# of a TypeError it swallowed; the guard reads the real per-engine floor now, so an
|
||||
# older engine with a genuine 32 floor still gets rejected with a clear message.
|
||||
MIN_SECONDS = 1
|
||||
MAX_SECONDS = SA3Inference.DIT_MAX_L * SAMPLES_PER_LATENT / SAMPLE_RATE // 1 # 380 s
|
||||
|
||||
def on_dit_change(dit_name):
|
||||
variants = discover_variants(dit_name)
|
||||
suggested = DEFAULT_DECODERS.get(dit_name, "same-s")
|
||||
dec_variants = discover_decoder_variants(suggested)
|
||||
var_u = (gr.update(choices=[(l, str(p)) for l, p in variants], value=str(variants[0][1]))
|
||||
if variants else gr.update(choices=[], value=None))
|
||||
dec_u = (gr.update(choices=[(l, str(p)) for l, p in dec_variants],
|
||||
value=str(dec_variants[0][1]))
|
||||
if dec_variants else gr.update(choices=[], value=None))
|
||||
return var_u, gr.update(value=suggested), dec_u
|
||||
|
||||
def on_decoder_change(decoder_name):
|
||||
dv = discover_decoder_variants(decoder_name)
|
||||
if not dv:
|
||||
return gr.update(choices=[], value=None)
|
||||
return gr.update(choices=[(l, str(p)) for l, p in dv], value=str(dv[0][1]))
|
||||
|
||||
def on_seconds_change(sec):
|
||||
"""Keep the inpaint range inside the clip length."""
|
||||
return gr.update(maximum=float(sec)), gr.update(maximum=float(sec))
|
||||
|
||||
# ── one generation, packaged as a history entry ──────────────────────────────────────────
|
||||
def _entry(dit_name, decoder_name, variant_path, dec_variant_path, prompt, negative_prompt,
|
||||
seconds, steps, seed_text, cfg, apg, sigma_max, chunking, fmt,
|
||||
a2a_path, a2a_sigma, inpaint_path, inp_start, inp_end, debug=False):
|
||||
# An empty prompt is a legitimate request, not an error: it generates
|
||||
# unconditionally. The encoder takes the empty string happily, and both
|
||||
# renderers already caption it "(no prompt)".
|
||||
prompt = (prompt or "").strip()
|
||||
try:
|
||||
seed = int(seed_text.strip()) if seed_text and seed_text.strip() else None
|
||||
except ValueError:
|
||||
return None, "seed must be an integer"
|
||||
if str(variant_path) == PT_EAGER_VARIANT:
|
||||
return None, "PT FP32 GT is not wired through this backend yet"
|
||||
try:
|
||||
inf = get_inference(dit_name, decoder_name, str(variant_path), str(dec_variant_path),
|
||||
default_T_lat, default_steps, default_seconds, quiet,
|
||||
chunking=bool(chunking))
|
||||
except Exception as e:
|
||||
return None, f"load failed: {type(e).__name__}: {e}"
|
||||
|
||||
# Which reference wins, and -- just as important -- SAY SO when something the user
|
||||
# supplied is being dropped. Silently falling through to a different mode is the
|
||||
# worst outcome here: upload an inpaint reference, forget the range sliders, and a
|
||||
# bare if/elif quietly runs audio-to-audio instead.
|
||||
#
|
||||
# The engines take ONE reference (generate_eager reads a single init_audio_path and
|
||||
# derives the inpaint keep-mask from its latents), so inpaint and audio-to-audio
|
||||
# cannot use two different files the way the MLX backend can.
|
||||
kw, mode, notes = {}, "text-to-audio", []
|
||||
sec = float(seconds)
|
||||
rng = None
|
||||
if inpaint_path and float(inp_end) > float(inp_start) and float(inp_start) < sec:
|
||||
if float(inp_end) > sec:
|
||||
notes.append(f"inpaint end clamped to the clip length ({sec:g}s)")
|
||||
rng = (float(inp_start), float(min(float(inp_end), sec)))
|
||||
elif inpaint_path and float(inp_end) <= float(inp_start):
|
||||
notes.append("inpainting ignored — set the start/end range sliders")
|
||||
elif inpaint_path:
|
||||
notes.append(f"inpainting ignored — start ({float(inp_start):g}s) is past "
|
||||
f"the end of the clip ({sec:g}s)")
|
||||
elif float(inp_end) > float(inp_start):
|
||||
notes.append("inpaint range ignored — no reference audio uploaded")
|
||||
|
||||
if rng is not None:
|
||||
kw = dict(init_audio_path=inpaint_path, inpaint_range=rng,
|
||||
init_noise_level=float(sigma_max))
|
||||
mode = "inpaint"
|
||||
if a2a_path:
|
||||
notes.append("using the inpaint reference — one reference per render, and "
|
||||
"inpainting needs it to keep the untouched part")
|
||||
elif a2a_path:
|
||||
# a2a's own slider IS the schedule start (parent-repo semantics), so it
|
||||
# overrides sigma_max rather than stacking with it.
|
||||
kw = dict(init_audio_path=a2a_path, init_noise_level=float(a2a_sigma))
|
||||
mode = "audio-to-audio"
|
||||
if float(sigma_max) != 1.0:
|
||||
notes.append(f"sigma_max {float(sigma_max):g} ignored — audio-to-audio uses "
|
||||
f"its own noise level ({float(a2a_sigma):g})")
|
||||
elif float(sigma_max) != 1.0:
|
||||
kw = dict(init_noise_level=float(sigma_max))
|
||||
neg = (negative_prompt or "").strip()
|
||||
if float(cfg) != 1.0:
|
||||
kw.update(cfg=float(cfg), apg=float(apg), negative_prompt=neg or None)
|
||||
elif neg:
|
||||
notes.append("negative prompt ignored — it only applies when CFG > 1")
|
||||
try:
|
||||
pcm, t = inf.generate(prompt, seconds=float(seconds), steps=int(steps),
|
||||
seed=seed, **kw)
|
||||
except NotImplementedError as e:
|
||||
return None, f"not supported by this engine: {e}"
|
||||
except Exception as e:
|
||||
return None, f"{type(e).__name__}: {e}"
|
||||
|
||||
debug_html = ""
|
||||
if debug:
|
||||
try:
|
||||
# Reuse the seed so the second pass is the same take, then hand both timing
|
||||
# dicts to the renderer. When the render already went eager there is nothing
|
||||
# to compare against, so skip the extra pass entirely.
|
||||
if not t.get("eager") and not getattr(inf, "_debug_warmed", False):
|
||||
# First eager call on this instance builds the _SecondContext set and its
|
||||
# buffers. Timing that pass charges ~600 ms of one-off setup to "unfused"
|
||||
# and makes the fusion win look far bigger than it is. Burn one pass.
|
||||
inf.generate_eager(prompt, seconds=float(seconds), steps=int(steps),
|
||||
seed=t["seed"], **kw)
|
||||
inf._debug_warmed = True
|
||||
t_stage = t if t.get("eager") else inf.generate_eager(
|
||||
prompt, seconds=float(seconds), steps=int(steps), seed=t["seed"], **kw)[1]
|
||||
debug_html = render_debug(inf, t, t_stage)
|
||||
except Exception as e:
|
||||
debug_html = (f"<div style='color:#e8a33d; font-size:0.82em'>debug pass failed: "
|
||||
f"{html_lib.escape(type(e).__name__)}: {html_lib.escape(str(e))}</div>")
|
||||
base = verbose_basename(prompt or "unconditional", neg,
|
||||
float(cfg), float(sigma_max), t["seed"])
|
||||
out_path = OUTPUT_DIR / f"{base}.wav"
|
||||
_save_wav(pcm, out_path)
|
||||
if fmt in ("flac", "mp3"):
|
||||
conv = out_path.with_suffix("." + fmt)
|
||||
rc = subprocess.run(["ffmpeg", "-y", "-loglevel", "error",
|
||||
"-i", str(out_path), str(conv)],
|
||||
capture_output=True).returncode
|
||||
if rc == 0 and conv.exists():
|
||||
out_path = conv
|
||||
mime = {"wav": "audio/wav", "flac": "audio/flac", "mp3": "audio/mpeg"}[out_path.suffix[1:]]
|
||||
try:
|
||||
t0 = time.time()
|
||||
png = render_spectrogram_png(pcm, sample_rate=SAMPLE_RATE, width=1400, height=280)
|
||||
spec_b64 = base64.b64encode(png).decode("ascii")
|
||||
spec_ms = (time.time() - t0) * 1000
|
||||
except Exception:
|
||||
spec_b64, spec_ms = "", 0.0
|
||||
timing = (f"engine-load {t.get('graph_build_ms', 0):.0f} ms · "
|
||||
f"Inference : {t['inference_ms']:.1f} ms · "
|
||||
f"{t.get('realtime', 0):.0f}× realtime · "
|
||||
f"seed : {t['seed']} · T_lat : {t['T_lat']} · "
|
||||
f"samples : {t['samples']} · spec {spec_ms:.0f} ms")
|
||||
return {
|
||||
"notes": notes, "debug_html": debug_html,
|
||||
"key": f"k{time.time_ns()}", "ts": time.time(), "dit": dit_name,
|
||||
"neg": neg, "cfg": float(cfg), "smx": float(sigma_max),
|
||||
"path": str(out_path), "mime": mime, "name": out_path.name,
|
||||
"size_mb": out_path.stat().st_size / 1e6, "mode": mode,
|
||||
"prompt": prompt, "seed": t["seed"],
|
||||
"lora": "", # kept for the shared renderer's contract
|
||||
"spec_b64": spec_b64, "timing": timing,
|
||||
}, None
|
||||
|
||||
def _present(state, entry, opts, queued_panel, *, force_autoplay=False):
|
||||
"""(player, timing, error, history, queued, state) for the shared renderers."""
|
||||
auto = force_autoplay or ("Auto-play" in (opts or []))
|
||||
radio = "Radio" in (opts or [])
|
||||
loop = "Loop" in (opts or [])
|
||||
hist = state.get("history", [])
|
||||
player = render_player(entry, autoplay=auto, radio=radio, advance=radio, loop=loop,
|
||||
autodl=("Auto-download" in (opts or []))) if entry else ""
|
||||
note_html = "".join(
|
||||
f"<div style='color:#e8a33d;font-size:0.85em'>note: {html_lib.escape(n)}</div>"
|
||||
for n in (entry.get("notes") or [])) if entry else ""
|
||||
return (player, (entry.get("timing", "") + note_html
|
||||
+ entry.get("debug_html", "")) if entry else "", "",
|
||||
render_history(hist, advance=radio, loop=loop),
|
||||
queued_panel, state)
|
||||
|
||||
CTRL = None # bound after the widgets exist
|
||||
|
||||
def _sig(ctrl):
|
||||
"""What the queued take was rendered FROM.
|
||||
|
||||
A pre-rendered take is only reusable while the controls still describe it. Serving
|
||||
it blind means the click after any edit silently returns the OLD settings: generate
|
||||
30 s, drag Seconds to 380, hit Generate, and you get 30 s of audio back. Everything
|
||||
in CTRL is a str/float/bool/None (audio inputs are filepaths), so str() over the
|
||||
tuple is a sound identity -- and it fails CLOSED, since anything it cannot compare
|
||||
simply looks different and forces a fresh render.
|
||||
"""
|
||||
return tuple(str(c) for c in ctrl)
|
||||
|
||||
_BLANK = {"current": None, "queued": None, "queued_sig": None, "history": []}
|
||||
|
||||
def generate(*args):
|
||||
*ctrl, opts, state = args
|
||||
state = dict(state or _BLANK)
|
||||
q = state.get("queued")
|
||||
if q is not None and state.get("queued_sig") == _sig(ctrl):
|
||||
# still the same request — show the pre-rendered take instantly
|
||||
state["queued"] = state["queued_sig"] = None
|
||||
state["current"] = q
|
||||
state["history"] = ([q] + state.get("history", []))[:24]
|
||||
return _present(state, q, opts, render_queue_status(generating=True),
|
||||
force_autoplay=True)
|
||||
# a control moved since it was queued: that take answers a question no longer
|
||||
# being asked, so drop it and render what the controls say now.
|
||||
state["queued"] = state["queued_sig"] = None
|
||||
entry, err = _entry(*ctrl)
|
||||
if err:
|
||||
return ("", "", f"<span style='color:#f88'>error: {html_lib.escape(err)}</span>",
|
||||
render_history(state.get("history", [])), render_queue_status(), state)
|
||||
state["current"] = entry
|
||||
state["history"] = ([entry] + state.get("history", []))[:24]
|
||||
return _present(state, entry, opts, render_queue_status(), force_autoplay=True)
|
||||
|
||||
def pregen(*args):
|
||||
"""Render the NEXT take in the background so the following click is instant."""
|
||||
*ctrl, opts, state = args
|
||||
state = dict(state or _BLANK)
|
||||
if ctrl[-1]:
|
||||
# Debug mode already renders twice (fused + unfused). Pre-rendering on top would
|
||||
# make it four renders per click for a take whose numbers may never be looked at.
|
||||
return "", state
|
||||
sig = _sig(ctrl)
|
||||
if state.get("queued") is not None and state.get("queued_sig") == sig:
|
||||
return render_queue_status(state["queued"]), state
|
||||
entry, err = _entry(*ctrl)
|
||||
state["queued"] = entry if not err else None
|
||||
state["queued_sig"] = sig if not err else None
|
||||
return render_queue_status(state.get("queued")), state
|
||||
|
||||
with gr.Blocks(title="SA3 · TensorRT", css=_css) as demo:
|
||||
st = gr.State(dict(_BLANK))
|
||||
with gr.Row():
|
||||
with gr.Column(scale=3):
|
||||
with gr.Row():
|
||||
dit_dd = gr.Dropdown(label="DiT model", choices=DIT_OPTIONS,
|
||||
value=initial_dit, scale=1)
|
||||
dec_dd = gr.Dropdown(label="Decoder", choices=["same-s", "same-l"],
|
||||
value=initial_decoder, scale=1)
|
||||
with gr.Row():
|
||||
var_dd = gr.Dropdown(label="DiT variant", allow_custom_value=True,
|
||||
choices=[(l, str(p)) for l, p in initial_variants],
|
||||
value=initial_variant_path, scale=1)
|
||||
dec_var_dd = gr.Dropdown(label="Decoder variant", allow_custom_value=True,
|
||||
choices=[(l, str(p)) for l, p in initial_dec_variants],
|
||||
value=initial_dec_variant_path, scale=1)
|
||||
with gr.Row():
|
||||
prompt = gr.Textbox(label="Prompt", lines=2, scale=6,
|
||||
placeholder="warm analog house groove")
|
||||
seed = gr.Textbox(label="Seed (optional)", max_lines=1, value="", scale=1)
|
||||
with gr.Row():
|
||||
seconds = gr.Slider(label="Seconds", minimum=MIN_SECONDS,
|
||||
maximum=MAX_SECONDS, step=1,
|
||||
value=default_seconds, scale=2)
|
||||
steps = gr.Slider(label="Steps", minimum=1, maximum=16, step=1,
|
||||
value=default_steps, scale=1)
|
||||
cfg = gr.Slider(label="CFG", minimum=0.0, maximum=10.0, step=0.1,
|
||||
value=1.0, scale=1)
|
||||
with gr.Accordion("Advanced", open=False):
|
||||
with gr.Row():
|
||||
apg = gr.Slider(label="APG (only applies when CFG > 1)",
|
||||
minimum=0.0, maximum=1.0, step=0.05, value=1.0)
|
||||
sigma_global = gr.Slider(label="sigma_max", minimum=0.05, maximum=1.0,
|
||||
step=0.05, value=1.0)
|
||||
negative_prompt = gr.Textbox(label="Negative prompt", lines=1)
|
||||
debug = gr.Checkbox(
|
||||
label="Debug: per-stage ms + VRAM", value=False,
|
||||
info="Adds a stage table under the timings. The mega-graph cannot be "
|
||||
"instrumented from inside (a captured graph rejects event records), "
|
||||
"so the split is measured by re-running the same engines unfused — "
|
||||
"which roughly doubles the render. The fused total is shown beside "
|
||||
"it, so the gap is the fusion win.")
|
||||
chunking = gr.Checkbox(
|
||||
label="Chunked decode (low VRAM)", value=True,
|
||||
info="On: 256-latent windows — 485 MB of decoder scratch on SAME-L, 346 MB "
|
||||
"on SAME-S. Off: single-shot on the wide profile — only 2-5% faster "
|
||||
"end-to-end (the DiT dominates), but 5-7 GB more resident. The scratch "
|
||||
"is a profile-ceiling reservation, so switching it off costs the same "
|
||||
"at 30 s as at 380 s.")
|
||||
with gr.Accordion("Audio-to-audio (guide the whole clip)", open=False):
|
||||
a2a_audio = gr.Audio(label="Guide audio — generation starts from its latents",
|
||||
type="filepath")
|
||||
a2a_sigma = gr.Slider(label="sigma_max (lower = closer to the guide)",
|
||||
minimum=0.05, maximum=1.0, step=0.05, value=0.6)
|
||||
with gr.Accordion("Inpainting (regenerate a span)", open=False):
|
||||
inpaint_audio = gr.Audio(label="Reference audio", type="filepath")
|
||||
with gr.Row():
|
||||
inp_start = gr.Slider(label="Start (s)", minimum=0,
|
||||
maximum=default_seconds, step=0.5, value=0)
|
||||
inp_end = gr.Slider(label="End (s)", minimum=0,
|
||||
maximum=default_seconds, step=0.5, value=0)
|
||||
with gr.Accordion("Output", open=False):
|
||||
output_opts = gr.CheckboxGroup(
|
||||
label="Playback", choices=["Auto-play", "Radio", "Loop", "Auto-download"],
|
||||
value=["Auto-play"])
|
||||
file_format = gr.Radio(label="File format", choices=["wav", "flac", "mp3"],
|
||||
value="wav")
|
||||
generate_btn = gr.Button("Generate", variant="primary", size="lg")
|
||||
promote_btn = gr.Button("", elem_id="sa3-promote") # CSS-hidden, DOM-present
|
||||
with gr.Column(scale=4):
|
||||
output_player = gr.HTML()
|
||||
timing = gr.HTML()
|
||||
error_box = gr.HTML()
|
||||
queued_html = gr.HTML()
|
||||
history_html = gr.HTML()
|
||||
|
||||
dit_dd.change(on_dit_change, inputs=[dit_dd], outputs=[var_dd, dec_dd, dec_var_dd])
|
||||
dec_dd.change(on_decoder_change, inputs=[dec_dd], outputs=[dec_var_dd])
|
||||
seconds.change(on_seconds_change, inputs=[seconds], outputs=[inp_start, inp_end]
|
||||
).then(None, js=_JS_FIX_SLIDERS)
|
||||
|
||||
CTRL = [dit_dd, dec_dd, var_dd, dec_var_dd, prompt, negative_prompt, seconds, steps,
|
||||
seed, cfg, apg, sigma_global, chunking, file_format,
|
||||
a2a_audio, a2a_sigma, inpaint_audio, inp_start, inp_end, debug]
|
||||
main_out = [output_player, timing, error_box, history_html, queued_html, st]
|
||||
# ⚠ js= takes a JS *function expression* ("() => ...") — gradio evaluates it as one.
|
||||
# _JS_TRY_PLAY is an inline-attribute body (bare statements, `this` = the <audio>),
|
||||
# so passing it here is a syntax error that kills the whole frontend: the config
|
||||
# fails to evaluate and the page sits on "Loading..." forever with zero controls.
|
||||
# The autoplay rescue belongs where it already is — render_player's oncanplay.
|
||||
for btn in (generate_btn, promote_btn):
|
||||
btn.click(generate, inputs=CTRL + [output_opts, st], outputs=main_out
|
||||
).then(pregen, inputs=CTRL + [output_opts, st],
|
||||
outputs=[queued_html, st]
|
||||
).then(None, js=_JS_FIX_SLIDERS)
|
||||
gr.Markdown("Renders are written to `output/`. **Radio** auto-advances to the next take as "
|
||||
"each finishes; the next one is pre-rendered while you listen.")
|
||||
|
||||
demo.queue(max_size=16).launch(share=share, server_name="0.0.0.0", show_error=True,
|
||||
allowed_paths=[str(OUTPUT_DIR)])
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--dit", choices=list(DEFAULT_DECODERS.keys()),
|
||||
default="sm-music",
|
||||
help="Initial DiT bundle (switchable at runtime)")
|
||||
ap.add_argument("--decoder", choices=["same-s", "same-l"], default=None,
|
||||
help="Initial decoder. Default: pairs with --dit")
|
||||
ap.add_argument("--default-seconds", type=float, default=120.0,
|
||||
help="Length to pre-warm the initial graph at")
|
||||
ap.add_argument("--default-steps", type=int, default=8)
|
||||
ap.add_argument("--share", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="Create a public gradio.live URL (default on)")
|
||||
ap.add_argument("--chunking", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="SAME-L: decode/encode in windows on the engines' low profile "
|
||||
"(509 MB of decoder scratch instead of 8143). --no-chunking uses the "
|
||||
"wide profile: single-shot, faster above L=256, ~5.4 GB more resident.")
|
||||
ap.add_argument("--quiet", action="store_true")
|
||||
args = ap.parse_args()
|
||||
global _CHUNKING
|
||||
_CHUNKING = args.chunking
|
||||
|
||||
if args.decoder is None:
|
||||
args.decoder = DEFAULT_DECODERS[args.dit]
|
||||
|
||||
T_lat = max(1, math.ceil(args.default_seconds * SAMPLE_RATE / SAMPLES_PER_LATENT))
|
||||
if args.decoder == "same-s" and T_lat % 2 != 0:
|
||||
T_lat += 1
|
||||
|
||||
print(f"\n━━━ SA3 TRT — gradio (debug build) ━━━")
|
||||
print(f" initial dit: {args.dit}")
|
||||
print(f" initial decoder: {args.decoder}")
|
||||
print(f" warmup: T_lat={T_lat} steps={args.default_steps} "
|
||||
f"(~{args.default_seconds}s)")
|
||||
# Pretty-print what variants are visible
|
||||
for ditname in DEFAULT_DECODERS:
|
||||
vs = discover_variants(ditname)
|
||||
print(f" dit-variants[{ditname}]: " + (
|
||||
", ".join(lbl for lbl, _ in vs) if vs else "(none found)"))
|
||||
for dec in ("same-s", "same-l"):
|
||||
dvs = discover_decoder_variants(dec)
|
||||
print(f" dec-variants[{dec}]: " + (
|
||||
", ".join(lbl for lbl, _ in dvs) if dvs else "(none found)"))
|
||||
print()
|
||||
|
||||
build_ui(args.dit, args.decoder, share=args.share, quiet=args.quiet,
|
||||
default_T_lat=T_lat, default_steps=args.default_steps,
|
||||
default_seconds=args.default_seconds)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,173 @@
|
||||
"""Mel-spectrogram renderer ported from Underfit's dashboard.
|
||||
|
||||
Source: /weka2/cj/clod/underfit/dashboard/server.py — the 3-band tinted stereo
|
||||
mel spectrogram used in the Underfit run dashboards. Same algorithm,
|
||||
same constants, same Slaney mel scale + librosa-compatible mel filterbank.
|
||||
|
||||
Differences vs the Underfit original:
|
||||
- Input is an int16 (T, 2) numpy array (in-memory PCM) rather than a file
|
||||
path; no soundfile dependency. SAMPLE_RATE assumed 44.1 kHz, resampled
|
||||
to 32 kHz internally for the spec (matches Underfit's target_sr).
|
||||
- Output is PNG bytes returned in-memory (no file write).
|
||||
- Exposes render_spectrogram_png(pcm_int16, sample_rate, width, height)
|
||||
so callers can choose a render resolution.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import io
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# Underfit's 3-band tint
|
||||
SPEC_BANDS = [
|
||||
(0, 200, (1.0, 0.0, 0.0)), # Bass -> Red
|
||||
(200, 1500, (0.0, 1.0, 0.0)), # Mid -> Green
|
||||
(1500, 16000, (0.0, 0.0, 1.0)), # High -> Blue
|
||||
]
|
||||
_BAND_COLORS = np.array([c for _, _, c in SPEC_BANDS], dtype=np.float32)
|
||||
|
||||
# Slaney mel-scale constants (linear below 1 kHz, log above), matches librosa default.
|
||||
_F_SP = 200.0 / 3
|
||||
_MIN_LOG_HZ = 1000.0
|
||||
_MIN_LOG_MEL = _MIN_LOG_HZ / _F_SP
|
||||
_LOGSTEP = math.log(6.4) / 27.0
|
||||
|
||||
|
||||
def _hz_to_mel(hz):
|
||||
hz = np.asarray(hz, dtype=np.float64)
|
||||
log_term = np.log(np.maximum(hz, _MIN_LOG_HZ) / _MIN_LOG_HZ) / _LOGSTEP
|
||||
return np.where(hz >= _MIN_LOG_HZ, _MIN_LOG_MEL + log_term, hz / _F_SP)
|
||||
|
||||
|
||||
def _mel_to_hz(mels):
|
||||
mels = np.asarray(mels, dtype=np.float64)
|
||||
return np.where(mels >= _MIN_LOG_MEL,
|
||||
_MIN_LOG_HZ * np.exp(_LOGSTEP * (mels - _MIN_LOG_MEL)),
|
||||
_F_SP * mels)
|
||||
|
||||
|
||||
def _mel_frequencies(n_mels, fmax, fmin=0.0):
|
||||
return _mel_to_hz(np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels))
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _mel_filterbank(n_mels, n_fft, sr, fmax):
|
||||
pts = _mel_to_hz(np.linspace(_hz_to_mel(0.0), _hz_to_mel(fmax), n_mels + 2))
|
||||
fft_f = np.linspace(0, sr / 2, n_fft // 2 + 1)
|
||||
filt = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
|
||||
for i in range(n_mels):
|
||||
lo, ce, hi = pts[i], pts[i + 1], pts[i + 2]
|
||||
left = (fft_f - lo) / max(ce - lo, 1e-10)
|
||||
right = (hi - fft_f) / max(hi - ce, 1e-10)
|
||||
filt[i] = np.maximum(0, np.minimum(left, right))
|
||||
enorm = (2.0 / (pts[2:n_mels + 2] - pts[0:n_mels])).astype(np.float32)
|
||||
filt *= enorm[:, None]
|
||||
return torch.from_numpy(filt)
|
||||
|
||||
|
||||
def _melspectrogram(y_ch, sr, n_mels=30, fmax=16000, hop_length=2048, n_fft=2048):
|
||||
y_t = torch.from_numpy(np.ascontiguousarray(y_ch)).float()
|
||||
win = torch.hann_window(n_fft)
|
||||
spec = torch.stft(y_t, n_fft=n_fft, hop_length=hop_length, window=win,
|
||||
center=True, return_complex=True, pad_mode='reflect')
|
||||
return (_mel_filterbank(n_mels, n_fft, sr, fmax) @ spec.abs().square()).numpy()
|
||||
|
||||
|
||||
def _power_to_db(S, top_db=80.0):
|
||||
log_spec = 10.0 * np.log10(np.maximum(S, 1e-10))
|
||||
return np.maximum(log_spec - log_spec.max(), -top_db)
|
||||
|
||||
|
||||
def _mel_channel(y_ch, sr, n_mels=30):
|
||||
"""dB-scaled mel + band-tinted RGB for one mono channel."""
|
||||
S = _melspectrogram(y_ch, sr, n_mels=n_mels, fmax=16000, hop_length=2048)
|
||||
|
||||
S_db = _power_to_db(S)
|
||||
np.clip(S_db, -60, 0, out=S_db)
|
||||
S_db += 60.0
|
||||
S_db /= 60.0
|
||||
np.power(S_db, 0.6, out=S_db)
|
||||
|
||||
mel_f = _mel_frequencies(n_mels, fmax=16000)
|
||||
n_frames = S.shape[1]
|
||||
band_norms = np.empty((3, n_frames), dtype=np.float32)
|
||||
for i, (flo, fhi, _) in enumerate(SPEC_BANDS):
|
||||
mask = (mel_f >= flo) & (mel_f < fhi)
|
||||
if mask.any():
|
||||
power = np.sum(S[mask], axis=0)
|
||||
db = 10.0 * np.log10(power + 1e-10)
|
||||
np.clip(db, -20, None, out=db)
|
||||
db -= -20
|
||||
mx = db.max()
|
||||
if mx > 0:
|
||||
db /= mx
|
||||
band_norms[i] = db
|
||||
else:
|
||||
band_norms[i] = 0.0
|
||||
|
||||
rgb = band_norms.T @ _BAND_COLORS
|
||||
for c in range(3):
|
||||
mx = rgb[:, c].max()
|
||||
if mx > 0:
|
||||
rgb[:, c] /= mx
|
||||
|
||||
return S_db, rgb
|
||||
|
||||
|
||||
def _resample_to_32k(pcm_int16, sr_in):
|
||||
"""Resample (T, 2) int16 PCM to (2, T') float32 at 32 kHz."""
|
||||
target_sr = 32000
|
||||
if sr_in == target_sr:
|
||||
y = pcm_int16.astype(np.float32) / 32768.0
|
||||
return y.T, target_sr
|
||||
# Convert int16 → float32 → (2, T), then torch interp.
|
||||
y = pcm_int16.astype(np.float32).T / 32768.0 # (2, T)
|
||||
y_t = torch.from_numpy(np.ascontiguousarray(y))
|
||||
new_len = int(round(y_t.shape[-1] * target_sr / sr_in))
|
||||
y_t = torch.nn.functional.interpolate(
|
||||
y_t.unsqueeze(0), size=new_len, mode='linear', align_corners=False
|
||||
).squeeze(0)
|
||||
return y_t.numpy(), target_sr
|
||||
|
||||
|
||||
def render_spectrogram_png(pcm_int16, sample_rate=44100, width=1200, height=240,
|
||||
n_mels=30) -> bytes:
|
||||
"""Render Underfit-style 3-band tinted stereo mel spectrogram to PNG bytes.
|
||||
|
||||
Args:
|
||||
pcm_int16: (T, 2) int16 numpy array — stereo PCM.
|
||||
sample_rate: input PCM sample rate (default 44100).
|
||||
width, height: final image size (default 1200x240; Underfit ships 300x60).
|
||||
n_mels: mel filterbank size (30 matches Underfit).
|
||||
|
||||
Returns: PNG bytes ready to base64 / save / etc.
|
||||
"""
|
||||
if pcm_int16.ndim != 2 or pcm_int16.shape[1] != 2:
|
||||
raise ValueError(f"expected (T, 2) PCM, got {pcm_int16.shape}")
|
||||
y, sr = _resample_to_32k(pcm_int16, sample_rate)
|
||||
# y is (2, T') float32
|
||||
S_L, rgb_L = _mel_channel(y[0], sr, n_mels=n_mels)
|
||||
S_R, rgb_R = _mel_channel(y[1], sr, n_mels=n_mels)
|
||||
|
||||
nf = min(S_L.shape[1], S_R.shape[1])
|
||||
S_L, S_R = S_L[:, :nf], S_R[:, :nf]
|
||||
rgb_L, rgb_R = rgb_L[:nf], rgb_R[:nf]
|
||||
nm = S_L.shape[0]
|
||||
|
||||
# L channel: flip vertically so bass-band is at bottom.
|
||||
S_L = S_L[::-1]
|
||||
|
||||
img = np.empty((nm * 2, nf, 3), dtype=np.float32)
|
||||
img[:nm] = S_L[:, :, np.newaxis] * rgb_L[np.newaxis, :, :]
|
||||
img[nm:] = S_R[:, :, np.newaxis] * rgb_R[np.newaxis, :, :]
|
||||
|
||||
np.clip(img, 0, 1, out=img)
|
||||
img *= 255
|
||||
pil = Image.fromarray(img.astype(np.uint8)).resize((width, height), Image.LANCZOS)
|
||||
buf = io.BytesIO()
|
||||
pil.save(buf, format="PNG", optimize=True)
|
||||
return buf.getvalue()
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Triton-compiled block-tiled MMA kernel, packaged as AOT PTX for trtp.aot_impl.
|
||||
|
||||
The SAME-L SWA plugin needs an ahead-of-time kernel so its TRT engine is
|
||||
CUDA-graph-capturable. The existing AOT kernel (diff_swa_ptx_kernel.py) is hand-written
|
||||
scalar-FP32 PTX — one warp per query, no K/V reuse, zero mma instructions — and is ~3.6x
|
||||
slower per layer than the Triton JIT path, which across the decoder's 12 attention
|
||||
layers accounts for the whole 55 -> 68 ms engine-level gap on sm_90.
|
||||
|
||||
This module closes that gap by compiling triton_diff_swa_mma's block-tiled kernel
|
||||
ahead of time and exposing its PTX. Using Triton as the code generator rather than
|
||||
hand-writing mma fragment layouts is deliberate: it already emits exactly the wanted
|
||||
codegen (verified: 512 x mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32), and PTX
|
||||
MMA fragment register assignment is error-prone to write by hand.
|
||||
|
||||
Two things have to be reconciled with what trtp.aot_impl can actually pass a kernel:
|
||||
|
||||
1. The `extra` channel carries INT32s only, so every non-int runtime scalar must be a
|
||||
constexpr. `scale` (fp32), `H_OUT`, `WINDOW`, the block sizes and 15 of the 16
|
||||
strides are therefore pinned at compile time, leaving only `N` at runtime.
|
||||
2. Triton appends two global-scratch pointers to the entry signature. TRT passes
|
||||
input/output pointers plus the extras and has no channel for them — but they are
|
||||
DECLARED AND NEVER LOADED (checked: only the declaration mentions them, there is no
|
||||
ld.param), so they are stripped from the PTX entry signature. That makes the
|
||||
kernel's declared parameter space exactly match what TRT supplies, rather than
|
||||
relying on the driver tolerating a short argument list.
|
||||
"""
|
||||
import re
|
||||
|
||||
import triton
|
||||
from triton.compiler.compiler import ASTSource
|
||||
|
||||
from triton_diff_swa_mma import (diff_swa_mma_kernel, WINDOW, BLOCK_N, BLOCK_D,
|
||||
BLOCK_KV)
|
||||
|
||||
HEAD_DIM = BLOCK_D
|
||||
_cache: dict = {}
|
||||
|
||||
|
||||
def _strip_unused_tail_params(ptx: str, keep: int) -> str:
|
||||
"""Remove entry parameters after the first `keep`, asserting they are unread.
|
||||
|
||||
Triton's trailing global-scratch pointers have no channel through aot_impl. They
|
||||
are never dereferenced, so dropping the declarations is safe — but verify that
|
||||
before cutting, so this fails loudly if a future Triton starts using them.
|
||||
"""
|
||||
m = re.search(r"(\.visible \.entry \w+\()([^)]*)(\))", ptx, re.S)
|
||||
if m is None:
|
||||
raise RuntimeError("could not locate the PTX entry signature")
|
||||
params = [p.strip() for p in m.group(2).split(",") if p.strip()]
|
||||
if len(params) <= keep:
|
||||
return ptx
|
||||
for p in params[keep:]:
|
||||
name = p.split()[-1]
|
||||
if re.search(rf"ld\.param[^\n]*\[{re.escape(name)}\]", ptx):
|
||||
raise RuntimeError(
|
||||
f"refusing to strip {name}: it IS loaded by the kernel, so TRT must "
|
||||
f"pass it and this packaging is invalid")
|
||||
new_sig = ",\n\t".join(params[:keep])
|
||||
return ptx[:m.start()] + m.group(1) + "\n\t" + new_sig + "\n" + m.group(3) + ptx[m.end():]
|
||||
|
||||
|
||||
def build(num_heads: int, arch: str):
|
||||
"""(kernel_name, ptx, shared_mem_bytes, num_warps) for this head count.
|
||||
|
||||
`arch` is the sm_XX being built for; the PTX `.target` is rewritten to it without
|
||||
the trailing 'a' (architecture-specific) suffix, which TRT's loader rejects. The
|
||||
m16n8k8 TF32 mma this kernel uses is available from sm_80, so the plain target is
|
||||
sufficient.
|
||||
"""
|
||||
key = (num_heads, arch)
|
||||
if key in _cache:
|
||||
return _cache[key]
|
||||
|
||||
h2 = 2 * num_heads
|
||||
# Order must match the kernel: inputs, then the runtime extra, then outputs.
|
||||
sig = {"Q": "*fp32", "K": "*fp32", "V": "*fp32", "N": "i32", "Out": "*fp32"}
|
||||
sig.update({f"stride_{t}{a}": "constexpr" for t in "qkvo" for a in "bnhd"})
|
||||
sig.update({"scale": "constexpr", "H_OUT": "constexpr",
|
||||
"WINDOW": "constexpr", "BLOCK_N": "constexpr",
|
||||
"BLOCK_D": "constexpr", "BLOCK_KV": "constexpr"})
|
||||
# Contiguous (B, N, 2H, D) in and (B, N, H, D) out. The batch strides are only ever
|
||||
# multiplied by program_id(2), which is 0 for the decoder's batch of 1, so pinning
|
||||
# them to 0 is exact here rather than an approximation.
|
||||
consts = {
|
||||
"stride_qb": 0, "stride_qn": h2 * HEAD_DIM, "stride_qh": HEAD_DIM, "stride_qd": 1,
|
||||
"stride_kb": 0, "stride_kn": h2 * HEAD_DIM, "stride_kh": HEAD_DIM, "stride_kd": 1,
|
||||
"stride_vb": 0, "stride_vn": h2 * HEAD_DIM, "stride_vh": HEAD_DIM, "stride_vd": 1,
|
||||
"stride_ob": 0, "stride_on": num_heads * HEAD_DIM, "stride_oh": HEAD_DIM,
|
||||
"stride_od": 1,
|
||||
"scale": HEAD_DIM ** -0.5, "H_OUT": num_heads, "WINDOW": WINDOW,
|
||||
"BLOCK_N": BLOCK_N, "BLOCK_D": BLOCK_D, "BLOCK_KV": BLOCK_KV,
|
||||
}
|
||||
try:
|
||||
src = ASTSource(diff_swa_mma_kernel, signature=sig, constexprs=consts)
|
||||
except TypeError: # older Triton spells it `constants`
|
||||
src = ASTSource(diff_swa_mma_kernel, signature=sig, constants=consts)
|
||||
|
||||
compiled = triton.compile(src)
|
||||
ptx = compiled.asm["ptx"]
|
||||
md = compiled.metadata
|
||||
|
||||
n_mma = len(re.findall(r"mma\.sync", ptx))
|
||||
if n_mma == 0:
|
||||
raise RuntimeError("compiled PTX has no mma instructions — tl.dot did not "
|
||||
"lower to tensor cores; the kernel would be no faster than "
|
||||
"the scalar hand-written one")
|
||||
|
||||
# TRT passes: Q, K, V (inputs), Out (output), then the extras. One extra (N).
|
||||
ptx = _strip_unused_tail_params(ptx, keep=5)
|
||||
ptx = re.sub(r"\.target\s+sm_\w+", f".target {arch}", ptx)
|
||||
# fp32 loads need 4-byte alignment; Triton emits .align 1 on the pointer params.
|
||||
ptx = ptx.replace(".align 1 .b8", ".align 4 .b8")
|
||||
|
||||
_cache[key] = (md.name, ptx, int(md.shared), int(md.num_warps))
|
||||
return _cache[key]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,170 @@
|
||||
"""Block-tiled differential SWA attention in Triton, shaped for AOT compilation.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
The SAME-L decoder's SWA plugin needs an ahead-of-time (PTX) implementation so its TRT
|
||||
engine is CUDA-graph-capturable — the JIT path re-enters Python per enqueue, which
|
||||
makes the engine non-capturable on sm_120 and silently drops the decode from the
|
||||
mega-graph. The existing AOT kernel (diff_swa_ptx_kernel.py) is hand-written PTX that
|
||||
is one warp per query, scalar FP32, no K/V reuse, and emits zero mma instructions —
|
||||
24% slower than the Triton JIT path on sm_90, and sweeping its only knob
|
||||
(warps_per_block over 1..16) moves it 7%, so the gap is algorithmic.
|
||||
|
||||
This module supplies the missing algorithm: BLOCK_N queries per block sharing the K/V
|
||||
window in SRAM, with tl.dot so QK^T and P·V land on tensor cores. Compiling it with
|
||||
triton.compile() yields PTX containing mma.sync instructions, which can then be handed
|
||||
to trtp.aot_impl — Triton's own codegen instead of hand-written fragment layouts.
|
||||
|
||||
The one substantive difference from triton_swa_v2.swa_attn_v2_kernel: that kernel
|
||||
computes all 2H heads and leaves the differential subtraction to the caller
|
||||
(o[:, :, :H] - o[:, :, H:]) in PyTorch. aot_impl returns a SINGLE kernel, so there is
|
||||
no caller to do that — the subtraction has to happen inside. Each block therefore owns
|
||||
one OUTPUT head h and processes two input heads, h (primary) and h + H (diff),
|
||||
subtracting before the store. That also saves materialising a 2H-head intermediate.
|
||||
|
||||
Precision note: tl.dot on fp32 inputs lowers to TF32 mma (10 mantissa bits), which is
|
||||
what the existing Triton path already does. The subtraction itself stays in FP32 on
|
||||
FP32 accumulator values, so the cancellation-sensitive step is unaffected — only the
|
||||
two attention products are TF32, exactly as in the shipped JIT path.
|
||||
"""
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def diff_swa_mma_kernel(
|
||||
# PARAMETER ORDER IS LOAD-BEARING: TRT's AOT plugin launcher passes
|
||||
# inputs..., extras..., outputs...
|
||||
# (confirmed against the working hand-written kernel in diff_swa_ptx_kernel.py,
|
||||
# whose entry is Q, K, V, N, 4 strides, Out — output LAST, after the extras).
|
||||
# An earlier version declared (Q, K, V, Out, N) and TRT therefore handed N where
|
||||
# the kernel expected the output pointer: it dereferenced 4352 as an address and
|
||||
# the engine died with "illegal memory access". Runtime scalars must sit between
|
||||
# the inputs and the outputs.
|
||||
Q, K, V,
|
||||
N,
|
||||
Out,
|
||||
stride_qb, stride_qn, stride_qh, stride_qd,
|
||||
stride_kb, stride_kn, stride_kh, stride_kd,
|
||||
stride_vb, stride_vn, stride_vh, stride_vd,
|
||||
stride_ob, stride_on, stride_oh, stride_od,
|
||||
scale, H_OUT,
|
||||
WINDOW: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_KV: tl.constexpr,
|
||||
):
|
||||
"""Differential SWA for one output head, one query block.
|
||||
|
||||
grid = (ceil(N / BLOCK_N), H_OUT, B)
|
||||
Out[b, n, h, :] = softmax(Qh Kh^T)Vh - softmax(Qh' Kh'^T)Vh' where h' = h + H_OUT
|
||||
"""
|
||||
block_n = tl.program_id(0)
|
||||
h = tl.program_id(1)
|
||||
b = tl.program_id(2)
|
||||
|
||||
q_start = block_n * BLOCK_N
|
||||
q_offs = tl.arange(0, BLOCK_N)
|
||||
q_pos = q_start + q_offs
|
||||
q_mask = q_pos < N
|
||||
|
||||
# K/V window spans [q_start - WINDOW, q_start + BLOCK_N + WINDOW)
|
||||
kv_start = q_start - WINDOW
|
||||
kv_offs = tl.arange(0, BLOCK_KV)
|
||||
kv_pos = kv_start + kv_offs
|
||||
kv_mask = (kv_pos >= 0) & (kv_pos < N)
|
||||
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
d_ok = d_offs[None, :] < BLOCK_D
|
||||
|
||||
# SWA mask is purely relative, so it is shared by both heads. kv_start is offset by
|
||||
# -WINDOW, hence |q_offs + WINDOW - kv_offs| <= WINDOW.
|
||||
rel = q_offs[:, None] + WINDOW - kv_offs[None, :]
|
||||
valid = (rel >= -WINDOW) & (rel <= WINDOW) & kv_mask[None, :] & q_mask[:, None]
|
||||
|
||||
acc = tl.zeros((BLOCK_N, BLOCK_D), dtype=tl.float32)
|
||||
|
||||
# head_sel 0 -> primary head h, 1 -> diff head h + H_OUT.
|
||||
#
|
||||
# tl.static_range, NOT a dynamic range(2). Measured: unrolled needs 64 KB of shared
|
||||
# memory (both heads' tiles live) and emits 512 mma; a dynamic loop is WORSE, not
|
||||
# better — Triton's pipeliner allocates multi-stage buffers for it and shared memory
|
||||
# jumps to 160 KB, with mma halved to 256. 64 KB exceeds the 48 KB *static* limit
|
||||
# and needs a cudaFuncSetAttribute opt-in from whoever launches it, which is the
|
||||
# open question for the TRT AOT path (see AOT_KERNEL_STATUS.md).
|
||||
for head_sel in tl.static_range(2):
|
||||
hh = h + head_sel * H_OUT
|
||||
|
||||
q_ptrs = (Q + b * stride_qb + q_pos[:, None] * stride_qn
|
||||
+ hh * stride_qh + d_offs[None, :] * stride_qd)
|
||||
q = tl.load(q_ptrs, mask=q_mask[:, None] & d_ok, other=0.0)
|
||||
|
||||
k_ptrs = (K + b * stride_kb + kv_pos[:, None] * stride_kn
|
||||
+ hh * stride_kh + d_offs[None, :] * stride_kd)
|
||||
k = tl.load(k_ptrs, mask=kv_mask[:, None] & d_ok, other=0.0)
|
||||
|
||||
scores = tl.dot(q, tl.trans(k)) * scale # tensor cores
|
||||
scores = tl.where(valid, scores, float("-inf"))
|
||||
|
||||
row_max = tl.max(scores, axis=1)
|
||||
e = tl.exp(scores - row_max[:, None])
|
||||
e = tl.where(valid, e, 0.0)
|
||||
row_sum = tl.sum(e, axis=1)
|
||||
w = e / row_sum[:, None]
|
||||
|
||||
v_ptrs = (V + b * stride_vb + kv_pos[:, None] * stride_vn
|
||||
+ hh * stride_vh + d_offs[None, :] * stride_vd)
|
||||
v = tl.load(v_ptrs, mask=kv_mask[:, None] & d_ok, other=0.0)
|
||||
|
||||
out_h = tl.dot(w, v) # tensor cores
|
||||
|
||||
# FP32 accumulator: primary minus diff. The cancellation-sensitive step stays
|
||||
# in FP32 even though the two products above are TF32.
|
||||
acc += tl.where(head_sel == 0, out_h, -out_h)
|
||||
|
||||
o_ptrs = (Out + b * stride_ob + q_pos[:, None] * stride_on
|
||||
+ h * stride_oh + d_offs[None, :] * stride_od)
|
||||
tl.store(o_ptrs, acc, mask=q_mask[:, None] & d_ok)
|
||||
|
||||
|
||||
# Launch geometry, kept here so the AOT path and the reference path agree exactly.
|
||||
#
|
||||
# BLOCK_KV MUST be >= BLOCK_N + 2*WINDOW or the K/V tile does not cover every position
|
||||
# the block's queries can attend to, and attention contributions are silently dropped.
|
||||
# (BLOCK_N=32 with BLOCK_KV=64 needs 66 and is therefore WRONG, even though it compiles
|
||||
# and fits in shared memory.)
|
||||
#
|
||||
# BLOCK_N=16 rather than 64, for two reasons measured on the real shapes:
|
||||
# speed 0.306 ms/layer vs 0.432 at BLOCK_N=64 (JIT reference: 0.421). With a window
|
||||
# of only +/-17, a 128-wide K/V tile is mostly masked-out waste; a 64-wide tile
|
||||
# keeps the useful fraction high.
|
||||
# memory 40 KB of shared memory vs 64 KB. 64 KB exceeds the 48 KB static limit and
|
||||
# needs a cudaFuncSetAttribute opt-in that the TRT AOT launcher does not do —
|
||||
# the BLOCK_N=64 engine built fine and then failed at enqueue with
|
||||
# "Failed to enqueue status -1", returning zeros.
|
||||
WINDOW = 17
|
||||
BLOCK_N = 16
|
||||
BLOCK_D = 64
|
||||
BLOCK_KV = 64 # >= BLOCK_N + 2 * WINDOW = 50
|
||||
|
||||
|
||||
def diff_swa_mma(q, k, v, window: int = WINDOW):
|
||||
"""Reference launcher. q, k, v: (B, N, 2H, D) fp32 -> (B, N, H, D) fp32.
|
||||
|
||||
Used to validate the kernel before it is compiled AOT; the shipped path drives the
|
||||
same kernel through trtp.aot_impl instead.
|
||||
"""
|
||||
B, Nn, H2, D = q.shape
|
||||
assert H2 % 2 == 0, f"expected an even head count, got {H2}"
|
||||
assert D == BLOCK_D, f"kernel is specialised for D={BLOCK_D}, got {D}"
|
||||
H = H2 // 2
|
||||
out = torch.empty((B, Nn, H, D), device=q.device, dtype=torch.float32)
|
||||
grid = (triton.cdiv(Nn, BLOCK_N), H, B)
|
||||
diff_swa_mma_kernel[grid](
|
||||
q, k, v, Nn, out,
|
||||
*q.stride(), *k.stride(), *v.stride(), *out.stride(),
|
||||
D ** -0.5, H,
|
||||
WINDOW=window, BLOCK_N=BLOCK_N, BLOCK_D=BLOCK_D, BLOCK_KV=BLOCK_KV,
|
||||
)
|
||||
return out
|
||||
@@ -0,0 +1,173 @@
|
||||
"""
|
||||
Block-tiled Triton SWA kernel using tl.dot for tensor core utilization.
|
||||
|
||||
Key optimization: process BLOCK_N queries per block, sharing K/V in SRAM.
|
||||
For SWA(17): window=35, so K/V per block = BLOCK_N + 34 positions.
|
||||
|
||||
Usage: CUDA_VISIBLE_DEVICES=5 python triton_swa_v2.py
|
||||
"""
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import time
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
|
||||
@triton.jit
|
||||
def swa_attn_v2_kernel(
|
||||
Q, K, V, Out,
|
||||
stride_qb, stride_qn, stride_qh, stride_qd,
|
||||
stride_kb, stride_kn, stride_kh, stride_kd,
|
||||
stride_vb, stride_vn, stride_vh, stride_vd,
|
||||
stride_ob, stride_on, stride_oh, stride_od,
|
||||
N, scale,
|
||||
WINDOW: tl.constexpr, # 17
|
||||
BLOCK_N: tl.constexpr, # queries per block (e.g., 64)
|
||||
BLOCK_D: tl.constexpr, # head dim (64)
|
||||
BLOCK_KV: tl.constexpr, # padded K/V window (next power of 2 >= BLOCK_N + 2*WINDOW)
|
||||
):
|
||||
"""Block-tiled SWA with online softmax."""
|
||||
block_n = tl.program_id(0) # Which query block
|
||||
h = tl.program_id(1) # Head
|
||||
b = tl.program_id(2) # Batch
|
||||
|
||||
# Query positions for this block
|
||||
q_start = block_n * BLOCK_N
|
||||
q_offs = tl.arange(0, BLOCK_N)
|
||||
q_pos = q_start + q_offs # (BLOCK_N,)
|
||||
q_mask = q_pos < N
|
||||
|
||||
# K/V window: [q_start - WINDOW, q_start + BLOCK_N + WINDOW)
|
||||
kv_start = q_start - WINDOW
|
||||
kv_offs = tl.arange(0, BLOCK_KV)
|
||||
kv_pos = kv_start + kv_offs # (BLOCK_KV,)
|
||||
kv_mask = (kv_pos >= 0) & (kv_pos < N)
|
||||
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
|
||||
# Load Q block: (BLOCK_N, BLOCK_D)
|
||||
q_ptrs = Q + b * stride_qb + q_pos[:, None] * stride_qn + h * stride_qh + d_offs[None, :] * stride_qd
|
||||
q = tl.load(q_ptrs, mask=q_mask[:, None] & (d_offs[None, :] < BLOCK_D), other=0.0).to(tl.float32)
|
||||
|
||||
# Load K window: (BLOCK_KV, BLOCK_D)
|
||||
k_ptrs = K + b * stride_kb + kv_pos[:, None] * stride_kn + h * stride_kh + d_offs[None, :] * stride_kd
|
||||
k = tl.load(k_ptrs, mask=kv_mask[:, None] & (d_offs[None, :] < BLOCK_D), other=0.0).to(tl.float32)
|
||||
|
||||
# Compute scores: Q @ K^T = (BLOCK_N, BLOCK_KV)
|
||||
scores = tl.dot(q, tl.trans(k)) * scale
|
||||
|
||||
# Apply SWA mask: query at q_pos[i] can only attend to kv_pos[j] where |q_pos[i] - kv_pos[j]| <= WINDOW
|
||||
# For each (i, j): check if |q_start + i - (kv_start + j)| <= WINDOW
|
||||
# = |i - j + WINDOW| <= WINDOW (since kv_start = q_start - WINDOW)
|
||||
# = |i - j + WINDOW| <= WINDOW
|
||||
# This simplifies to: j >= i and j < i + 2*WINDOW + 1
|
||||
# i.e., j in [i, i + 2*WINDOW + 1)
|
||||
rel_pos = q_offs[:, None] + WINDOW - kv_offs[None, :] # (BLOCK_N, BLOCK_KV)
|
||||
swa_mask = (rel_pos >= -WINDOW) & (rel_pos <= WINDOW)
|
||||
|
||||
# Combined mask: SWA + valid positions
|
||||
valid_mask = swa_mask & kv_mask[None, :] & q_mask[:, None]
|
||||
scores = tl.where(valid_mask, scores, float('-inf'))
|
||||
|
||||
# Row-wise softmax
|
||||
row_max = tl.max(scores, axis=1) # (BLOCK_N,)
|
||||
scores = scores - row_max[:, None]
|
||||
exp_scores = tl.exp(scores)
|
||||
exp_scores = tl.where(valid_mask, exp_scores, 0.0)
|
||||
row_sum = tl.sum(exp_scores, axis=1) # (BLOCK_N,)
|
||||
weights = exp_scores / row_sum[:, None]
|
||||
|
||||
# Load V window: (BLOCK_KV, BLOCK_D)
|
||||
v_ptrs = V + b * stride_vb + kv_pos[:, None] * stride_vn + h * stride_vh + d_offs[None, :] * stride_vd
|
||||
v = tl.load(v_ptrs, mask=kv_mask[:, None] & (d_offs[None, :] < BLOCK_D), other=0.0).to(tl.float32)
|
||||
|
||||
# Output: weights @ V = (BLOCK_N, BLOCK_D)
|
||||
out = tl.dot(weights.to(v.dtype), v)
|
||||
|
||||
# Store
|
||||
o_ptrs = Out + b * stride_ob + q_pos[:, None] * stride_on + h * stride_oh + d_offs[None, :] * stride_od
|
||||
tl.store(o_ptrs, out.to(tl.bfloat16), mask=q_mask[:, None] & (d_offs[None, :] < BLOCK_D))
|
||||
|
||||
|
||||
def triton_swa_attn_v2(q, k, v, window=17):
|
||||
"""Block-tiled Triton SWA. q,k,v: (B, N, H, D). Returns (B, N, H, D)."""
|
||||
B, N, H, D = q.shape
|
||||
out = torch.empty_like(q)
|
||||
scale = D ** -0.5
|
||||
|
||||
BLOCK_N = 64
|
||||
# K/V window per block = BLOCK_N + 2*window = 64 + 34 = 98 → pad to 128
|
||||
BLOCK_KV = 128 # Next power of 2 >= 98
|
||||
|
||||
num_blocks_n = (N + BLOCK_N - 1) // BLOCK_N
|
||||
grid = (num_blocks_n, H, B)
|
||||
|
||||
swa_attn_v2_kernel[grid](
|
||||
q, k, v, out,
|
||||
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
|
||||
k.stride(0), k.stride(1), k.stride(2), k.stride(3),
|
||||
v.stride(0), v.stride(1), v.stride(2), v.stride(3),
|
||||
out.stride(0), out.stride(1), out.stride(2), out.stride(3),
|
||||
N, scale,
|
||||
WINDOW=window,
|
||||
BLOCK_N=BLOCK_N,
|
||||
BLOCK_D=D,
|
||||
BLOCK_KV=BLOCK_KV,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from flash_attn import flash_attn_func
|
||||
|
||||
device = "cuda"
|
||||
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
||||
|
||||
# Test accuracy
|
||||
for B, N, H, D in [(1, 200, 48, 64), (1, 1836, 48, 64), (1, 5000, 48, 64)]:
|
||||
q = torch.randn(B, N, H, D, device=device, dtype=torch.bfloat16)
|
||||
k = torch.randn(B, N, H, D, device=device, dtype=torch.bfloat16)
|
||||
v = torch.randn(B, N, H, D, device=device, dtype=torch.bfloat16)
|
||||
|
||||
ref = flash_attn_func(q, k, v, window_size=(17, 17))
|
||||
out = triton_swa_attn_v2(q, k, v, window=17)
|
||||
|
||||
cos = F.cosine_similarity(ref.flatten().float(), out.flatten().float(), dim=0).item()
|
||||
print(f"B={B} N={N:5d} H={H}: cos={cos:.6f}")
|
||||
|
||||
# Benchmark
|
||||
print("\n--- Benchmark ---")
|
||||
for N_tok in [1836, 21964]:
|
||||
B, H, D = 1, 48, 64
|
||||
q = torch.randn(B, N_tok, H, D, device=device, dtype=torch.bfloat16)
|
||||
k = torch.randn(B, N_tok, H, D, device=device, dtype=torch.bfloat16)
|
||||
v = torch.randn(B, N_tok, H, D, device=device, dtype=torch.bfloat16)
|
||||
|
||||
for _ in range(5):
|
||||
triton_swa_attn_v2(q, k, v)
|
||||
flash_attn_func(q, k, v, window_size=(17, 17))
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times = []
|
||||
for _ in range(20):
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
triton_swa_attn_v2(q, k, v)
|
||||
torch.cuda.synchronize()
|
||||
times.append(time.perf_counter() - t0)
|
||||
ms_t = np.median(times) * 1000
|
||||
|
||||
times = []
|
||||
for _ in range(20):
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
flash_attn_func(q, k, v, window_size=(17, 17))
|
||||
torch.cuda.synchronize()
|
||||
times.append(time.perf_counter() - t0)
|
||||
ms_f = np.median(times) * 1000
|
||||
|
||||
ratio = ms_t / ms_f
|
||||
print(f"N={N_tok:6d}: v2={ms_t:.2f}ms flash={ms_f:.2f}ms ratio={ratio:.1f}x")
|
||||
print(f" ×12 layers: v2={ms_t*12:.1f}ms flash={ms_f*12:.1f}ms")
|
||||
@@ -0,0 +1,202 @@
|
||||
#!/usr/bin/env python3
|
||||
"""EngineInspector identity for the fp8-linears + native-FMHA + fp32-RoPE DiT engine.
|
||||
From the DETAILED per-layer I/O Format/Datatype (baked into the serialized engine), NOT
|
||||
the build flag. Three make-or-break checks the task demands:
|
||||
|
||||
CHECK 1 fp8 GEMMs fire on the Linears — count gemm-type layers running FP8 E4M3 (~222).
|
||||
CHECK 2 native fused MHA fires — count _gemm_mha_v2 / fmha nodes (expect 96,
|
||||
bf16 or fp16 — NOT 0/unfused, NOT fp8).
|
||||
CHECK 3 RoPE is BAKED, not computed at runtime — at most 2 cos/sin layers survive
|
||||
(timestep_features + seconds_total), and some
|
||||
fp32 remains. Boundary dtypes are reported but
|
||||
NOT asserted: they differ per arch by Myelin
|
||||
fusion (see the comment on check3).
|
||||
|
||||
usage: python verify_fp8_rope.py <engine.trt> [label]
|
||||
"""
|
||||
import sys, json, re
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
import tensorrt as trt
|
||||
|
||||
|
||||
def norm_dtype(s):
|
||||
if not s:
|
||||
return None
|
||||
t = str(s).lower()
|
||||
if "int8" in t: return "INT8"
|
||||
if "fp8" in t or "e4m3" in t or "e5m2" in t: return "FP8"
|
||||
if "bf16" in t or "bfloat" in t: return "BF16"
|
||||
if "fp16" in t or "half" in t or "float16" in t: return "FP16"
|
||||
if "fp32" in t or "float32" in t or "float " in t or t.endswith("float"): return "FP32"
|
||||
return None
|
||||
|
||||
|
||||
def io_dtypes(ly):
|
||||
found = set(); raw = []
|
||||
for io_key in ("Inputs", "Outputs"):
|
||||
for io in ly.get(io_key, []) or []:
|
||||
if isinstance(io, dict):
|
||||
for fk in ("Format/Datatype", "Datatype", "Format", "Dtype"):
|
||||
v = io.get(fk)
|
||||
if v:
|
||||
raw.append(str(v)); n = norm_dtype(v)
|
||||
if n: found.add(n)
|
||||
for key in ("Precision", "ComputePrecision"):
|
||||
v = ly.get(key)
|
||||
if v:
|
||||
raw.append(f"{key}={v}"); n = norm_dtype(v)
|
||||
if n: found.add(n)
|
||||
return found, raw
|
||||
|
||||
|
||||
def main():
|
||||
path = sys.argv[1]
|
||||
label = sys.argv[2] if len(sys.argv) > 2 else Path(path).name
|
||||
logger = trt.Logger(trt.Logger.ERROR)
|
||||
runtime = trt.Runtime(logger)
|
||||
engine = runtime.deserialize_cuda_engine(Path(path).read_bytes())
|
||||
if engine is None:
|
||||
raise SystemExit(f"deserialize failed: {path}")
|
||||
insp = engine.create_engine_inspector()
|
||||
info = json.loads(insp.get_engine_information(trt.LayerInformationFormat.JSON))
|
||||
layers = info.get("Layers", info if isinstance(info, list) else [])
|
||||
|
||||
print(f"\n{'='*80}\nENGINE: {label}\n path: {path}\n"
|
||||
f" size: {Path(path).stat().st_size/1e9:.2f} GB\n total layers: {len(layers)}")
|
||||
|
||||
# An engine built without ProfilingVerbosity.DETAILED serializes layer NAMES only:
|
||||
# the inspector hands back a list of strings, every per-layer lookup finds nothing,
|
||||
# and all three checks report False — which reads exactly like "fp8 did not fire"
|
||||
# when the truth is "nothing was recorded to look at". Say so instead of guessing.
|
||||
n_dict = sum(1 for ly in layers if isinstance(ly, dict))
|
||||
if layers and n_dict == 0:
|
||||
print(f"\n {'!'*76}")
|
||||
print(" CANNOT VERIFY: this engine was built without ProfilingVerbosity.DETAILED,")
|
||||
print(" so it carries layer names only — no per-layer Format/Datatype to inspect.")
|
||||
print(" The checks below will all report False regardless of what the engine does.")
|
||||
print(" Rebuild for inspection with:")
|
||||
print(" cfg = builder.create_builder_config()")
|
||||
print(" cfg.profiling_verbosity = trt.ProfilingVerbosity.DETAILED")
|
||||
print(" Verbosity is metadata-only — it changes neither kernel selection nor")
|
||||
print(" numerics (measured: identical velocity-cos, speed within 0.7%), so the")
|
||||
print(" DETAILED twin is a valid stand-in for the engine you ship.")
|
||||
print(f" {'!'*76}")
|
||||
|
||||
type_hist = Counter(); gemm_prec = Counter(); prec_any = Counter()
|
||||
gemm_examples = []; mha_nodes = []; mha_examples = []; fp8_layers = 0
|
||||
rope_fp32 = []; rope_seen = []
|
||||
MHA_RE = re.compile(r"mha|fmha", re.I)
|
||||
# timestep_features and seconds_total are the only Fourier chains evaluated at
|
||||
# runtime; RoPE's own cos/sin are baked constants and must not appear as layers.
|
||||
MAX_RUNTIME_COSSIN = 2
|
||||
ROPE_RE = re.compile(r"(cos|sin|q_norm|k_norm|rotary|inv_freq|angle)", re.I)
|
||||
|
||||
for ly in layers:
|
||||
if not isinstance(ly, dict):
|
||||
continue
|
||||
name = ly.get("Name", "?")
|
||||
lt = str(ly.get("LayerType", ly.get("ParameterType", "?")))
|
||||
type_hist[lt] += 1
|
||||
found, raw = io_dtypes(ly)
|
||||
for p in found:
|
||||
prec_any[p] += 1
|
||||
if "FP8" in found:
|
||||
fp8_layers += 1
|
||||
if lt == "gemm":
|
||||
dom = ("FP8" if "FP8" in found else "INT8" if "INT8" in found else
|
||||
"FP16" if "FP16" in found else "BF16" if "BF16" in found else
|
||||
"FP32" if "FP32" in found else "UNKNOWN")
|
||||
gemm_prec[dom] += 1
|
||||
if dom == "FP8" and len(gemm_examples) < 4:
|
||||
gemm_examples.append((name, raw[:4]))
|
||||
blob = json.dumps(ly)
|
||||
if MHA_RE.search(name) or MHA_RE.search(blob):
|
||||
mha_nodes.append((name, lt, sorted(found)))
|
||||
if len(mha_examples) < 4:
|
||||
mha_examples.append((name, lt, raw[:6]))
|
||||
if ROPE_RE.search(name) and ("cos" in name.lower() or "sin" in name.lower()
|
||||
or "q_norm" in name.lower() or "k_norm" in name.lower()):
|
||||
rope_seen.append(name)
|
||||
if "FP32" in found:
|
||||
rope_fp32.append(name)
|
||||
|
||||
n_gemm = sum(gemm_prec.values()); n_gemm_fp8 = gemm_prec.get("FP8", 0)
|
||||
mha_prec = Counter()
|
||||
for _, _, found in mha_nodes:
|
||||
d = ("FP8" if "FP8" in found else "BF16" if "BF16" in found else
|
||||
"FP16" if "FP16" in found else "FP32" if "FP32" in found else "OTHER")
|
||||
mha_prec[d] += 1
|
||||
|
||||
print("\n--- layer-type histogram (top 14) ---")
|
||||
for lt, c in type_hist.most_common(14):
|
||||
print(f" {c:6d} {lt}")
|
||||
|
||||
print(f"\n{'#'*80}\nCHECK 1 — FP8 GEMMs FIRE ON LINEARS?")
|
||||
for p in ("FP8", "INT8", "FP16", "BF16", "FP32", "UNKNOWN"):
|
||||
if gemm_prec.get(p):
|
||||
print(f" {p:8s}: {gemm_prec[p]:5d}")
|
||||
print(f" => {n_gemm_fp8}/{n_gemm} gemm layers run FP8 (E4M3). total FP8-carrying layers: {fp8_layers}")
|
||||
for nm, raw in gemm_examples[:3]:
|
||||
print(f" {nm[:78]}"); [print(f" {r[:86]}") for r in raw]
|
||||
# native-FMHA fp8 count matches the fp8_gemm reference (176), NOT the sage-plugin
|
||||
# variant (222): with native bf16-FMHA, TRT fuses attn-adjacent linears differently.
|
||||
check1 = n_gemm_fp8 >= 170
|
||||
|
||||
print(f"\n{'#'*80}\nCHECK 2 — NATIVE FUSED MHA FIRES (NOT 0/unfused)?")
|
||||
print(f" fused-MHA (_gemm_mha_v2/fmha) nodes: {len(mha_nodes)} precision: {dict(mha_prec)}")
|
||||
for nm, lt, raw in mha_examples:
|
||||
print(f" ({lt}) {nm[:74]}"); [print(f" {r[:86]}") for r in raw]
|
||||
check2 = len(mha_nodes) >= 90 and mha_prec.get("FP8", 0) == 0
|
||||
|
||||
# CHECK 3 rewritten 2026-07-31. The old condition was
|
||||
# len(rope_fp32) > 0 and prec_any["FP32"] > 100
|
||||
# and it reported False on BOTH architectures for the engine it exists to validate.
|
||||
# Two reasons, both about what this metadata can actually tell you:
|
||||
#
|
||||
# 1. ">100 FP32 layers" describes the fp32-TRUNK recipe. The shipped engine BAKES
|
||||
# RoPE's cos/sin as fp32 constants, so there is no fp32 trunk — only the two
|
||||
# runtime Fourier chains (timestep_features, seconds_total) stay fp32, which is
|
||||
# 6 FP32 layers on sm_120 and 16 on sm_90. Neither will ever exceed 100.
|
||||
#
|
||||
# 2. The per-layer "FP32" here is the dtype at the FUSED BLOB BOUNDARY, not the
|
||||
# precision of the arithmetic inside it, and Myelin fuses those boundaries
|
||||
# differently per arch. Measured on the same ONNX and recipe:
|
||||
# sm_120 __myl_MaxMinDivReshMulSinCastCosCastConc I/O ['Float','BFloat16']
|
||||
# sm_90 __myl_CastMaxMinDivReshMulCosSinCastCastConc I/O ['BFloat16','BFloat16']
|
||||
# Same engine identity, same measured fidelity, opposite verdicts from a
|
||||
# name-and-boundary check. So "is the island fp32 inside" is NOT decidable here.
|
||||
#
|
||||
# What IS decidable and does matter: RoPE must be BAKED, i.e. no runtime cos/sin
|
||||
# beyond those two Fourier chains. If the bake silently failed, RoPE would be
|
||||
# computed at runtime in reduced precision and drift at long sequence — the bug this
|
||||
# recipe exists to avoid. Fidelity itself belongs to a numerical test, not an
|
||||
# inspector: sweep velocity-cos vs the fp32 engine across t (calibration and RoPE
|
||||
# damage both show up at t~1.0, and are invisible mid-schedule).
|
||||
print(f"\n{'#'*80}\nCHECK 3 — RoPE BAKED (no runtime cos/sin beyond the 2 Fourier chains)?")
|
||||
print(f" cos/sin-named layers: {len(rope_seen)} (expect <= {MAX_RUNTIME_COSSIN}: "
|
||||
f"timestep_features + seconds_total)")
|
||||
for nm in rope_seen[:6]:
|
||||
print(f" {'FP32' if nm in rope_fp32 else 'reduced-precision boundary'}: {nm}")
|
||||
print(f" precision histogram across ALL layers: {dict(prec_any)}")
|
||||
print(f" note: boundary dtypes above are informational — they differ by arch "
|
||||
f"(Myelin fusion), so they do not decide the island. Use a velocity-cos sweep.")
|
||||
check3 = len(rope_seen) <= MAX_RUNTIME_COSSIN and prec_any.get("FP32", 0) > 0
|
||||
|
||||
out = dict(engine=path, label=label, total_layers=len(layers),
|
||||
gemm_by_precision=dict(gemm_prec), n_gemm=n_gemm, n_gemm_fp8=n_gemm_fp8,
|
||||
total_fp8_layers=fp8_layers, precision_any=dict(prec_any),
|
||||
n_mha_fused=len(mha_nodes), mha_precision=dict(mha_prec),
|
||||
n_rope_fp32=len(rope_fp32),
|
||||
check1_fp8_gemms=bool(check1),
|
||||
check2_native_mha_fused=bool(check2),
|
||||
n_runtime_cossin=len(rope_seen),
|
||||
check3_rope_baked=bool(check3))
|
||||
Path(path).with_suffix(".verify.json").write_text(json.dumps(out, indent=2))
|
||||
print(f"\n SUMMARY: CHECK1(fp8 gemms>=170)={check1} "
|
||||
f"CHECK2(mha fused, no fp8)={check2} CHECK3(rope baked)={check3}")
|
||||
print(f" wrote {Path(path).with_suffix('.verify.json')}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user