| """Whisper base building blocks for recipes/whisper-base.py. |
| |
| Everything here is built from the pinned openai/whisper-base checkpoint loaded through transformers |
| (WhisperForConditionalGeneration); the program wrappers reuse its nn.Modules (weights are shared, not |
| copied) and only re-express the forward pass so that it lowers to one WebGPU delegate per program. |
| |
| Program contract (same split and I/O as the playground's previous whisper-base export): |
| |
| mel forward(frames[1,3000,400] f32) -> power_mel[1,3000,80] f32 |
| frames: 30 s of 16 kHz audio, zero-padded to 480000 samples, reflect-padded by 200 on |
| both ends, 3000 windows of 400 samples at hop 160 (= whisper.js buildFrames). |
| power_mel = |rfft(hann_periodic * frame)|^2 @ mel_filters (no log; see host_log_mel). |
| encoder forward(log_mel[1,80,3000] f32) -> encoder_hidden[1,1500,512] f32 (fp32 weights) |
| cross_kv forward(encoder_hidden[1,1500,512]) -> 12 x [1,8,1500,64] f32, order k0,v0,...,k5,v5 |
| (per decoder layer: encoder_attn.k_proj / v_proj, heads-first [B,H,S,D]) |
| decoder_split forward(tok_embeds[1,64,512], k0, v0, ..., k5, v5) -> hidden[1,64,512] |
| tok_embeds: token embeddings WITHOUT positions for positions 0..64 (zero rows beyond |
| the live prefix); the program adds positions 0..63, runs causal self-attention over |
| the 64 rows and returns final-LayerNorm hidden states (row p predicts token p+1). |
| decode_stream forward(embed[1,1,512], input_pos[1] i64, k0, v0, ..., k5, v5) -> hidden[1,1,512] |
| embed = token embedding + positional embedding of `input_pos` (added on the host); |
| self-attention K/V cache [1,448,8,64] x 6 layers x {k,v} lives inside the program |
| (mutable buffers owned by the delegate); feed positions 0,1,2,... in order. |
| lmhead forward(hidden[1,1,512]) -> logits[1,1,51872] (tied to decoder.embed_tokens; rows |
| 51865..51871 are padding with logit -1e9, see LMHeadProgram) |
| |
| Precision (whisper-base.py --precision): |
| q4 (previous-export parity) cross_kv / decoder_split / decode_stream / lmhead linears are 4-bit groupwise |
| weight-only (torchao IntxWeightOnlyConfig int4, group 64, symmetric, "affine" qparams -> |
| et_vk.linear_q4gsw); int8 token-embedding table. mel and encoder are fp32. |
| fp32 every program fp32 (aten.linear on 2-D views); fp32 token-embedding table. |
| |
| Host-side files (reproduced from the checkpoint by the recipe, consumed by the playground JS): |
| whisper_embed.i8.bin int8 [51865,512] round(embed_tokens / scale), scale = rowwise absmax/127 (q4) |
| whisper_embed.scale.f32.bin f32 [51865] the per-row scales (q4) |
| whisper_embed.f32.bin f32 [51865,512] embed_tokens (fp32) |
| whisper_pos_embed.f32.bin f32 [448,512] decoder.embed_positions.weight |
| vocab.json upstream tokenizer vocab (byte-level BPE, detokenization only) |
| """ |
|
|
| import copy |
| import io |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| SAMPLE_RATE = 16000 |
| N_FFT = 400 |
| HOP = 160 |
| FRAMES = 3000 |
| N_SAMPLES = FRAMES * HOP |
| PAD = N_FFT // 2 |
| N_FREQ = N_FFT // 2 + 1 |
| N_MELS = 80 |
| ENC_LEN = 1500 |
| D_MODEL = 512 |
| N_HEADS = 8 |
| HEAD_DIM = 64 |
| N_LAYERS = 6 |
| MAX_DEC = 64 |
| KV_LEN = 448 |
| VOCAB = 51865 |
| VOCAB_PADDED = 51872 |
| PAD_LOGIT = -1e9 |
|
|
| TOK = { |
| "sot": 50258, |
| "eot": 50257, |
| "en": 50259, |
| "transcribe": 50359, |
| "notimestamps": 50363, |
| } |
| PREFIX_EN = [TOK["sot"], TOK["en"], TOK["transcribe"], TOK["notimestamps"]] |
|
|
|
|
| |
| |
| |
|
|
|
|
| def build_frames(pcm): |
| """Mirror of whisper.js buildFrames(): [3000, 400] float32 windows (no window function).""" |
| pcm = np.asarray(pcm, dtype=np.float32) |
| base = np.zeros(N_SAMPLES + 2 * PAD, dtype=np.float32) |
| n = min(pcm.size, N_SAMPLES) |
| base[PAD : PAD + n] = pcm[:n] |
| for i in range(PAD): |
| base[PAD - 1 - i] = base[PAD + 1 + i] |
| base[PAD + N_SAMPLES + i] = base[PAD + N_SAMPLES - 2 - i] |
| idx = np.arange(FRAMES)[:, None] * HOP + np.arange(N_FFT)[None, :] |
| return base[idx] |
|
|
|
|
| def host_log_mel(power_mel): |
| """Host step between the mel and encoder programs (WhisperFeatureExtractor normalization). |
| |
| power_mel [1,3000,80] (or [3000,80]) -> log_mel [1,80,3000]: |
| x = log10(max(p, 1e-10)); x = max(x, max(x) - 8); x = (x + 4) / 4; transpose. |
| float32 throughout (the JS equivalent uses Float32Array + Math.log10). |
| """ |
| p = np.asarray(power_mel, dtype=np.float32).reshape(FRAMES, N_MELS) |
| x = np.log10(np.maximum(p, np.float32(1e-10))).astype(np.float32) |
| x = np.maximum(x, x.max() - np.float32(8.0)) |
| x = ((x + np.float32(4.0)) / np.float32(4.0)).astype(np.float32) |
| return np.ascontiguousarray(x.T)[None] |
|
|
|
|
| class HostTables: |
| """Token-embedding lookup as the playground JS does it, plus the files that carry the tables. |
| |
| kind="int8": whisper_embed.i8.bin (int8 [51865,512], round(w / s)) + whisper_embed.scale.f32.bin |
| (s = rowwise absmax / 127): byte-identical to the previous export's files. |
| kind="fp32": whisper_embed.f32.bin (fp32 [51865,512], the checkpoint's embed_tokens). |
| Both: whisper_pos_embed.f32.bin (fp32 [448,512], decoder.embed_positions).""" |
|
|
| def __init__(self, model, kind): |
| e = model.model.decoder.embed_tokens.weight.detach().float() |
| self.kind = kind |
| self.pos = model.model.decoder.embed_positions.weight.detach().float().numpy() |
| if kind == "int8": |
| scale = e.abs().amax(dim=1) / 127.0 |
| self.table = torch.round(e / scale[:, None]).clamp(-127, 127).to(torch.int8).numpy() |
| self.scale = scale.numpy().astype(np.float32) |
| elif kind == "fp32": |
| self.table = e.numpy().astype(np.float32) |
| self.scale = None |
| else: |
| raise ValueError(kind) |
|
|
| def embed(self, ids, position=None): |
| e = self.table[ids].astype(np.float32) |
| if self.scale is not None: |
| e = e * self.scale[ids][..., None] |
| if position is not None: |
| e = e + self.pos[position] |
| return e.astype(np.float32) |
|
|
| def files(self): |
| out = {} |
| if self.kind == "int8": |
| out["whisper_embed.i8.bin"] = self.table.tobytes() |
| out["whisper_embed.scale.f32.bin"] = self.scale.astype("<f4").tobytes() |
| else: |
| out["whisper_embed.f32.bin"] = self.table.astype("<f4").tobytes() |
| out["whisper_pos_embed.f32.bin"] = self.pos.astype("<f4").tobytes() |
| return out |
|
|
|
|
| |
| |
| |
|
|
|
|
| def dft_mel_weights(feature_extractor): |
| """cos / -sin DFT rows with the periodic Hann window folded in, plus the HF mel filters.""" |
| n = np.arange(N_FFT) |
| k = np.arange(N_FREQ) |
| window = 0.5 - 0.5 * np.cos( |
| 2.0 * np.pi * n / N_FFT |
| ) |
| ang = 2.0 * np.pi * ((k[:, None] * n[None, :]) % N_FFT) / N_FFT |
| wc = (np.cos(ang) * window[None, :]).astype(np.float32) |
| ws = (-np.sin(ang) * window[None, :]).astype(np.float32) |
| mel = np.asarray(feature_extractor.mel_filters, dtype=np.float64).T.astype( |
| np.float32 |
| ) |
| assert mel.shape == (N_MELS, N_FREQ), mel.shape |
| return torch.from_numpy(wc), torch.from_numpy(ws), torch.from_numpy(mel) |
|
|
|
|
| class MelProgram(torch.nn.Module): |
| def __init__(self, feature_extractor): |
| super().__init__() |
| wc, ws, mel = dft_mel_weights(feature_extractor) |
| self.register_buffer("wc", wc) |
| self.register_buffer("ws", ws) |
| self.register_buffer("mel", mel) |
|
|
| def forward(self, frames): |
| x = frames.reshape( |
| FRAMES, N_FFT |
| ) |
| re = F.linear(x, self.wc) |
| im = F.linear(x, self.ws) |
| return F.linear(re * re + im * im, self.mel).reshape(1, FRAMES, N_MELS) |
|
|
|
|
| def _heads(x, s): |
| """[s, 512] or [1, s, 512] -> [1, 8, s, 64].""" |
| return x.reshape(1, s, N_HEADS, HEAD_DIM).transpose(1, 2) |
|
|
|
|
| def _merge(o, s): |
| """[1, 8, s, 64] -> [1, s, 512].""" |
| return o.transpose(1, 2).reshape(1, s, D_MODEL) |
|
|
|
|
| def attention(q, k, v, mask, impl): |
| """q/k/v [1,H,S,D]; scale = 1/sqrt(64) = 0.125 (a power of two, so applying it to q as HF does |
| or to q.k as the kernels do is bit-identical).""" |
| if impl == "et_vk": |
| import executorch.backends.vulkan.custom_ops_lib |
|
|
| return torch.ops.et_vk.sdpa(q, k, v, mask, None) |
| s = torch.matmul(q * 0.125, k.transpose(-2, -1)) |
| if mask is not None: |
| s = s + mask |
| return torch.matmul(torch.softmax(s, dim=-1), v) |
|
|
|
|
| class EncoderProgram(torch.nn.Module): |
| """HF WhisperEncoder forward re-expressed with 2-D fp32 linears. |
| |
| attn="eager": matmul + softmax (lowered to aten.bmm/_softmax; activation memory is planned). |
| attn="et_vk": the fused et_vk.sdpa op (what the previous export used); on the public WebGPU |
| runtime each call allocates its own two [8,1500,1500] f32 scratch buffers (72 MB each). |
| """ |
|
|
| def __init__(self, model, attn="eager"): |
| super().__init__() |
| self.enc = model.model.encoder |
| self.attn = attn |
|
|
| def forward(self, log_mel): |
| e = self.enc |
| x = F.gelu(e.conv1(log_mel)) |
| x = F.gelu(e.conv2(x)) |
| x = x.permute(0, 2, 1).reshape(ENC_LEN, D_MODEL) + e.embed_positions.weight |
| for layer in e.layers: |
| a = layer.self_attn |
| h = layer.self_attn_layer_norm(x) |
| q = _heads(a.q_proj(h), ENC_LEN) |
| k = _heads(a.k_proj(h), ENC_LEN) |
| v = _heads(a.v_proj(h), ENC_LEN) |
| o = attention(q, k, v, None, "et_vk" if self.attn == "et_vk" else "torch") |
| x = x + a.out_proj(o.transpose(1, 2).reshape(ENC_LEN, D_MODEL)) |
| h = layer.final_layer_norm(x) |
| x = x + layer.fc2(F.gelu(layer.fc1(h))) |
| return e.layer_norm(x).reshape(1, ENC_LEN, D_MODEL) |
|
|
|
|
| def _lin(m, x): |
| """nn.Linear that lowers on the public WebGPU runtime: fp32 aten.linear only accepts 2-D |
| inputs, so fp32 weights get a [-1, K] view; int4 (q4gsw) linears accept any rank and keep it.""" |
| if type(m.weight.data) is torch.Tensor and x.dim() != 2: |
| return m(x.reshape(-1, x.shape[-1])).reshape(*x.shape[:-1], m.out_features) |
| return m(x) |
|
|
|
|
| class CrossKVProgram(torch.nn.Module): |
| def __init__(self, decoder): |
| super().__init__() |
| self.layers = decoder.layers |
|
|
| def forward(self, enc): |
| out = [] |
| for layer in self.layers: |
| a = layer.encoder_attn |
| out.append(_heads(_lin(a.k_proj, enc), ENC_LEN)) |
| out.append(_heads(_lin(a.v_proj, enc), ENC_LEN)) |
| return tuple(out) |
|
|
|
|
| def _mlp(layer, x): |
| return x + _lin(layer.fc2, F.gelu(_lin(layer.fc1, layer.final_layer_norm(x)))) |
|
|
|
|
| def _cross(layer, x, k, v, s, impl): |
| a = layer.encoder_attn |
| h = layer.encoder_attn_layer_norm(x) |
| o = attention(_heads(_lin(a.q_proj, h), s), k, v, None, impl) |
| return x + _lin(a.out_proj, _merge(o, s)) |
|
|
|
|
| class DecoderSplitProgram(torch.nn.Module): |
| """Full-recompute decoder over MAX_DEC rows with a causal mask (no KV cache).""" |
|
|
| def __init__(self, decoder, attn="et_vk"): |
| super().__init__() |
| self.dec = decoder |
| self.attn = attn |
| self.register_buffer( |
| "pos", decoder.embed_positions.weight.detach()[:MAX_DEC].clone() |
| ) |
| causal = torch.full((MAX_DEC, MAX_DEC), -1e30).triu(1) |
| |
| self.register_buffer( |
| "mask", causal.expand(1, N_HEADS, MAX_DEC, MAX_DEC).contiguous() |
| ) |
|
|
| def forward(self, tok_embeds, *cross_kv): |
| x = tok_embeds + self.pos |
| for i, layer in enumerate(self.dec.layers): |
| a = layer.self_attn |
| h = layer.self_attn_layer_norm(x) |
| q = _heads(_lin(a.q_proj, h), MAX_DEC) |
| k = _heads(_lin(a.k_proj, h), MAX_DEC) |
| v = _heads(_lin(a.v_proj, h), MAX_DEC) |
| o = attention(q, k, v, self.mask, self.attn) |
| x = x + _lin(a.out_proj, _merge(o, MAX_DEC)) |
| x = _cross( |
| layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], MAX_DEC, self.attn |
| ) |
| x = _mlp(layer, x) |
| return self.dec.layer_norm(x) |
|
|
|
|
| class DecodeStreamProgram(torch.nn.Module): |
| """One token per call; self-attention K/V cache [1,448,8,64] per layer kept in the program. |
| |
| Self-attention uses llama.update_cache + llama.custom_sdpa (executorch.extension.llm.custom_ops), |
| which the Vulkan partitioner fuses into sdpa_with_kv_cache with an et_vk.select_as_symint |
| position. Cross-attention uses et_vk.sdpa (attn="et_vk") or matmul/softmax (attn="torch"). |
| """ |
|
|
| def __init__(self, decoder, attn="et_vk"): |
| super().__init__() |
| from executorch.extension.llm.custom_ops import custom_ops |
|
|
| self.dec = decoder |
| self.attn = attn |
| for i in range(N_LAYERS): |
| self.register_buffer( |
| f"k_cache_{i}", |
| torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM), |
| persistent=False, |
| ) |
| self.register_buffer( |
| f"v_cache_{i}", |
| torch.zeros(1, KV_LEN, N_HEADS, HEAD_DIM), |
| persistent=False, |
| ) |
|
|
| def reset(self): |
| for i in range(N_LAYERS): |
| getattr(self, f"k_cache_{i}").zero_() |
| getattr(self, f"v_cache_{i}").zero_() |
|
|
| def forward(self, embed, input_pos, *cross_kv): |
| pos = input_pos[0].item() |
| torch._check_is_size(pos) |
| torch._check(pos < KV_LEN) |
| x = embed |
| for i, layer in enumerate(self.dec.layers): |
| a = layer.self_attn |
| h = layer.self_attn_layer_norm(x) |
| q = _lin(a.q_proj, h).view(1, 1, N_HEADS, HEAD_DIM) |
| k = _lin(a.k_proj, h).view(1, 1, N_HEADS, HEAD_DIM) |
| v = _lin(a.v_proj, h).view(1, 1, N_HEADS, HEAD_DIM) |
| kc = getattr(self, f"k_cache_{i}") |
| vc = getattr(self, f"v_cache_{i}") |
| torch.ops.llama.update_cache(k, kc, pos) |
| torch.ops.llama.update_cache(v, vc, pos) |
| o = torch.ops.llama.custom_sdpa(q, kc, vc, pos, None, 0.0, True, None) |
| x = x + _lin(a.out_proj, o.reshape(1, 1, D_MODEL)) |
| x = _cross(layer, x, cross_kv[2 * i], cross_kv[2 * i + 1], 1, self.attn) |
| x = _mlp(layer, x) |
| return self.dec.layer_norm(x) |
|
|
|
|
| class LMHeadProgram(torch.nn.Module): |
| """Tied output projection. The vocab is padded to `pad_to` rows: executorch 1.5.1's q4gsw |
| lowering pads the scales to a multiple of 4 output channels and then fails its own shape |
| check for N = 51865, so N must already be a multiple of 4 (the previous export used 51872). |
| Padded rows have zero weights and a -1e9 bias, so they can never win an argmax.""" |
|
|
| def __init__(self, model, pad_to=0): |
| super().__init__() |
| w = model.model.decoder.embed_tokens.weight.detach().clone() |
| n = w.shape[0] |
| b = torch.zeros(max(n, pad_to)) |
| if pad_to and pad_to > n: |
| w = torch.cat([w, w.new_zeros(pad_to - n, w.shape[1])]) |
| b[n:] = PAD_LOGIT |
| self.proj = torch.nn.Linear(D_MODEL, w.shape[0], bias=True) |
| self.proj.weight = torch.nn.Parameter(w, requires_grad=False) |
| self.proj.bias = torch.nn.Parameter(b, requires_grad=False) |
|
|
| def forward(self, hidden): |
| return _lin(self.proj, hidden) |
|
|
|
|
| |
| |
| |
|
|
| Q4_DESC = ( |
| "torchao.quantization.quant_api.IntxWeightOnlyConfig(weight_dtype=torch.int4, " |
| "granularity=PerGroup(64), intx_choose_qparams_algorithm='affine') applied with quantize_ to " |
| "every nn.Linear of the program; symmetric, scale = group absmax / 7.5, q in [-8, 7]" |
| ) |
|
|
|
|
| def quantize_q4(module, group_size=64): |
| from torchao.quantization.granularity import PerGroup |
| from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ |
|
|
| quantize_( |
| module, |
| IntxWeightOnlyConfig( |
| weight_dtype=torch.int4, |
| granularity=PerGroup(group_size), |
| intx_choose_qparams_algorithm="affine", |
| ), |
| ) |
| return module |
|
|
|
|
| Q8_HEAD_DESC = ( |
| "PT2E with executorch.backends.vulkan.quantizer.VulkanQuantizer: weights int8 per-output-channel " |
| "symmetric (PerChannelMinMaxObserver, the quantizer's own weight spec), input activation int8 " |
| "per-tensor affine STATIC (HistogramObserver) calibrated on lmhead inputs; lowers to " |
| "quantized_decomposed.quantize_per_tensor + et_vk.linear_q8ta_q8csw" |
| ) |
|
|
|
|
| def q8_lmhead(lm, calib_hidden): |
| """Static-activation int8 x per-channel int8 lmhead (the only 8-bit linear executorch 1.5.1's |
| Vulkan lowering emits that the public WebGPU runtime registers). Returns the converted GraphModule |
| and the (input scale, zero point, dequantized weight) needed for an equivalent portable module.""" |
| from executorch.backends.vulkan.quantizer.vulkan_quantizer import ( |
| get_symmetric_quantization_config, |
| VulkanQuantizer, |
| ) |
| from torchao.quantization.pt2e import HistogramObserver |
| from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e |
| from torchao.quantization.pt2e.quantizer import QuantizationConfig, QuantizationSpec |
|
|
| base = get_symmetric_quantization_config(is_dynamic=False, weight_bits=8) |
| act = QuantizationSpec( |
| dtype=torch.int8, |
| quant_min=-128, |
| quant_max=127, |
| qscheme=torch.per_tensor_affine, |
| is_dynamic=False, |
| observer_or_fake_quant_ctr=HistogramObserver, |
| ) |
| cfg = QuantizationConfig( |
| input_activation=act, output_activation=None, weight=base.weight, bias=None, is_qat=False |
| ) |
| with torch.no_grad(): |
| gm = torch.export.export(lm.eval(), (calib_hidden[:1],), strict=False).module() |
| gm = prepare_pt2e(gm, VulkanQuantizer().set_global(cfg)) |
| for i in range(calib_hidden.shape[0]): |
| gm(calib_hidden[i : i + 1]) |
| gm = convert_pt2e(gm) |
| qp, wdq = None, None |
| for n in gm.graph.nodes: |
| if n.op != "call_function": |
| continue |
| name = str(n.target) |
| if "quantize_per_tensor" in name and "dequantize" not in name and qp is None: |
| qp = (float(n.args[1]), int(n.args[2])) |
| if "dequantize_per_channel" in name: |
| w = getattr(gm, n.args[0].target) |
| sc = getattr(gm, n.args[1].target) |
| wdq = (w.float() * sc.float()[:, None]).contiguous() |
| if qp is None or wdq is None: |
| raise RuntimeError("q8 lmhead: could not find the quantize/dequantize nodes") |
| return gm, qp, wdq |
|
|
|
|
| class FakeQ8LMHead(torch.nn.Module): |
| """Portable-runtime equivalent of the q8 lmhead: per-tensor int8 fake-quant of the input, then |
| an fp32 linear on the dequantized int8 weights (same bias).""" |
|
|
| def __init__(self, lm, qp, wdq): |
| super().__init__() |
| self.scale, self.zp = qp |
| self.register_buffer("w", wdq) |
| self.register_buffer("b", lm.proj.bias.detach().clone()) |
|
|
| def forward(self, hidden): |
| x = hidden.reshape(-1, D_MODEL) |
| q = torch.clamp(torch.round(x / self.scale) + self.zp, -128, 127) |
| x = (q - self.zp) * self.scale |
| return F.linear(x, self.w, self.b).reshape(*hidden.shape[:-1], self.w.shape[0]) |
|
|
|
|
| def word_errors(tokenizer, hyp_text, ref_text): |
| """(edits, reference words) after WhisperTokenizer.normalize (English text normalizer).""" |
| h = tokenizer.normalize(hyp_text).split() |
| r = tokenizer.normalize(ref_text).split() |
| d = list(range(len(r) + 1)) |
| for i in range(1, len(h) + 1): |
| prev, d[0] = d[0], i |
| for j in range(1, len(r) + 1): |
| cur = d[j] |
| d[j] = min(d[j] + 1, d[j - 1] + 1, prev + (h[i - 1] != r[j - 1])) |
| prev = cur |
| return d[len(r)], len(r) |
|
|
|
|
| def dequantized_copy(module): |
| """Same module with every quantized linear weight replaced by its fp32 dequantization. |
| |
| Used only for the CPU portable .pte (the portable runtime has no torchao kernels); its numerics |
| are the q4 numerics (dequantize-then-fp32-matmul, as the q4gsw kernels compute).""" |
| m = copy.deepcopy(module) |
| for sub in m.modules(): |
| if isinstance(sub, torch.nn.Linear) and hasattr(sub.weight, "dequantize"): |
| sub.weight = torch.nn.Parameter( |
| sub.weight.dequantize().float(), requires_grad=False |
| ) |
| return m |
|
|
|
|
| |
| |
| |
|
|
|
|
| def load_librispeech_dummy(parquet_path, max_sec=15.0, count=8, longer=False): |
| """First `count` clips (in file order) no longer than max_sec seconds, as 16 kHz float32. |
| longer=True selects the clips *longer* than max_sec instead (used only for calibration, so |
| calibration and evaluation clips never overlap).""" |
| import pyarrow.parquet as pq |
| import soundfile as sf |
|
|
| table = pq.read_table(parquet_path).to_pylist() |
| clips = [] |
| for row in table: |
| audio, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32") |
| if sr != SAMPLE_RATE: |
| raise ValueError(f"unexpected sample rate {sr}") |
| if audio.ndim != 1: |
| audio = audio.mean(axis=1) |
| if (audio.size / SAMPLE_RATE > max_sec) != longer: |
| continue |
| clips.append({"id": row["id"], "text": row["text"], "pcm": audio}) |
| if len(clips) >= count: |
| break |
| return clips |
|
|
|
|
| def seconds(pcm): |
| return round(len(pcm) / SAMPLE_RATE, 3) |
|
|
|
|
| def isclose_report(a, b): |
| a = np.asarray(a, dtype=np.float64).ravel() |
| b = np.asarray(b, dtype=np.float64).ravel() |
| d = np.abs(a - b) |
| cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30)) |
| return { |
| "max_abs": float(d.max()), |
| "mean_abs": float(d.mean()), |
| "max_rel_to_absmax": float(d.max() / (np.abs(b).max() + 1e-30)), |
| "cosine": cos, |
| } |
|
|
|
|
|
|
| |
| |
| |
|
|
|
|
| def _vk_graphs(pte_bytes): |
| from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( |
| extract_vk_flatbuffer, |
| flatbuffer_to_vk_graph, |
| ) |
| from executorch.exir._serialize._program import deserialize_pte_binary |
|
|
| pte = deserialize_pte_binary(pte_bytes) |
| prog = pte.program if hasattr(pte, "program") else pte |
| seg_base = int.from_bytes(pte_bytes[24:32], "little") |
| out = [] |
| for plan in prog.execution_plan: |
| for d in plan.delegates: |
| proc = d.processed |
| if proc.location.name == "INLINE": |
| blob = prog.backend_delegate_data[proc.index].data |
| else: |
| seg = prog.segments[proc.index] |
| blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size] |
| out.append(flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))) |
| return prog, out |
|
|
|
|
| def program_io(pte_bytes): |
| """Method names and their input/output tensor shapes + dtypes from the serialized program.""" |
| from executorch.exir._serialize._program import deserialize_pte_binary |
|
|
| pte = deserialize_pte_binary(pte_bytes) |
| prog = pte.program if hasattr(pte, "program") else pte |
| res = {} |
| for plan in prog.execution_plan: |
|
|
| def desc(i): |
| v = plan.values[i].val |
| if hasattr(v, "sizes"): |
| return {"shape": list(v.sizes), "dtype": v.scalar_type.name.lower()} |
| return {"kind": type(v).__name__} |
|
|
| res[plan.name] = { |
| "inputs": [desc(i) for i in plan.inputs], |
| "outputs": [desc(i) for i in plan.outputs], |
| } |
| return res |
|
|
|
|
| def vk_checks(pte_bytes, runtime_ops=None): |
| """Checks against constraints of the public WebGPU runtime's op implementations |
| (backends/webgpu/runtime/ops): op registered, aten.linear 2-D, et_vk.sdpa mask shape, |
| sdpa_with_kv_cache without mask / with causal, fp32 tensors only.""" |
| _, graphs = _vk_graphs(pte_bytes) |
| problems, notes = [], {} |
| for g in graphs: |
| vals = [v.value for v in g.values] |
|
|
| def t(i): |
| v = vals[i] |
| return v if type(v).__name__ == "VkTensor" else None |
|
|
| dtypes = sorted({type(v).__name__ if type(v).__name__ != "VkTensor" else v.datatype.name for v in vals}) |
| notes["value_kinds"] = dtypes |
| for op in g.chain: |
| name = op.name |
| if runtime_ops is not None and name not in runtime_ops: |
| problems.append(f"{name}: not registered in the WebGPU runtime op list") |
| a = op.args |
| if name == "aten.linear.default": |
| ri, rw = len(t(a[0]).dims), len(t(a[1]).dims) |
| notes.setdefault("aten.linear input ranks", set()).add(ri) |
| if ri != 2 or rw != 2: |
| problems.append(f"aten.linear with input rank {ri}, weight rank {rw} (runtime needs 2-D)") |
| if name == "et_vk.sdpa.default": |
| q, k = t(a[0]), t(a[1]) |
| m = t(a[3]) |
| if m is not None: |
| need = q.dims[0] * q.dims[1] * q.dims[2] * k.dims[2] |
| got = int(np.prod(m.dims)) |
| if need != got: |
| problems.append(f"et_vk.sdpa mask numel {got} != B*H*Sq*Skv {need}") |
| notes.setdefault("et_vk.sdpa q/k/mask", set()).add( |
| (tuple(q.dims), tuple(k.dims), tuple(m.dims) if m is not None else None) |
| ) |
| if name == "sdpa_with_kv_cache.default": |
| mask_kind = type(vals[a[7]]).__name__ |
| causal = vals[a[9]] |
| if mask_kind != "Null": |
| problems.append("sdpa_with_kv_cache with attn_mask (runtime rejects)") |
| if not (type(causal).__name__ == "Bool" and causal.bool_val): |
| problems.append(f"sdpa_with_kv_cache is_causal={causal} (runtime needs true)") |
| notes.setdefault("sdpa_with_kv_cache q/k_cache", set()).add( |
| (tuple(t(a[0]).dims), tuple(t(a[3]).dims), type(vals[a[5]]).__name__) |
| ) |
| if name == "aten.convolution.default": |
| notes.setdefault("aten.convolution input/weight", set()).add( |
| (tuple(t(a[0]).dims), tuple(t(a[1]).dims)) |
| ) |
| for v in vals: |
| if type(v).__name__ == "VkTensor" and v.datatype.name not in ("FLOAT32", "INT32", "INT64", "UINT8", "INT8"): |
| problems.append(f"tensor dtype {v.datatype.name}") |
| break |
| for k, v in list(notes.items()): |
| if isinstance(v, set): |
| notes[k] = sorted(str(x) for x in v) |
| return {"problems": problems, "ok": not problems, "notes": notes} |
|
|
