From 6c3a23e054477a39e2238ef8d2ea8e0f2436bae1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 22 Jul 2026 18:45:14 +0900 Subject: [PATCH] feat(xla): add prefill embeddings entry Emit a distinct StableHLO prefill path for post-scale embeddings with an explicit additive attention bias, stable input validation, frozen structural coverage, and real IREE parity checks for tied and untied checkpoints. Refs #858 --- spike/openxla/prefill_embeddings_check.py | 234 +++++++++++ src/lib/mlxcel-xla/README.md | 33 ++ .../prefill_embeddings_logits.mlir | 380 ++++++++++++++++++ .../mlxcel-xla/src/emitter/context_tests.rs | 24 +- src/lib/mlxcel-xla/src/emitter/mod.rs | 169 +++++++- src/lib/mlxcel-xla/src/emitter/model.rs | 331 +++++++++++++-- src/lib/mlxcel-xla/src/validation.rs | 35 +- 7 files changed, 1160 insertions(+), 46 deletions(-) create mode 100644 spike/openxla/prefill_embeddings_check.py create mode 100644 src/lib/mlxcel-xla/assets/qwen2-moe-tiny/prefill_embeddings_logits.mlir diff --git a/spike/openxla/prefill_embeddings_check.py b/spike/openxla/prefill_embeddings_check.py new file mode 100644 index 00000000..0748f5e8 --- /dev/null +++ b/spike/openxla/prefill_embeddings_check.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +"""Real-IREE parity check for token and embeddings StableHLO prefill entries. + +The fixture builds deterministic tiny Llama-shaped tied and untied checkpoints, +emits both prefill modules from each exact config, compiles them for IREE's +llvm-cpu/local-task target, and invokes them with one shared weight set. The +embeddings input is gathered from the checkpoint's token embedding table; the +explicit attention bias uses the emitter's finite 0/-1e30 causal convention. + +Every output is compared, not only the selected logits: both K and V tensors for +every layer and every padded bucket row must agree within ATOL. Scenarios cover a +one-token prompt, nonzero padding tokens, an untied LM head, and a prompt one token +below the 256-token bucket capacity. + +Run from the repository root with the OpenXLA spike environment: + + spike/openxla/.venv/bin/python spike/openxla/prefill_embeddings_check.py + +Exit 0 means both modules compiled and every parity scenario passed. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +from iree.compiler.tools import compile_file +from iree.runtime import load_vm_flatbuffer_file + +REPO_ROOT = Path(__file__).resolve().parents[2] +CARGO = os.environ.get("CARGO", "cargo") +PREFILL_LP = 256 +HIDDEN = 8 +INTERMEDIATE = 16 +N_LAYERS = 2 +N_Q = 2 +N_KV = 1 +HEAD_DIM = 4 +VOCAB = 32 +MASKED = np.float32(-1e30) +ATOL = 1e-5 + + +def emitter_config(*, tied: bool) -> dict[str, object]: + return { + "model_type": "llama", + "hidden_size": HIDDEN, + "intermediate_size": INTERMEDIATE, + "num_hidden_layers": N_LAYERS, + "num_attention_heads": N_Q, + "num_key_value_heads": N_KV, + "head_dim": HEAD_DIM, + "vocab_size": VOCAB, + "rms_norm_eps": 1e-6, + "rope_theta": 10_000.0, + "tie_word_embeddings": tied, + } + + +def random_weight(rng: np.random.Generator, shape: tuple[int, ...]) -> np.ndarray: + return np.ascontiguousarray(rng.normal(0.0, 0.08, shape), dtype=np.float32) + + +def checkpoint_weights(*, tied: bool) -> list[np.ndarray]: + """Return one synthetic checkpoint in the emitter's deterministic arg order.""" + rng = np.random.default_rng(858 if tied else 859) + weights = [ + random_weight(rng, (VOCAB, HIDDEN)), + np.ascontiguousarray(rng.uniform(0.8, 1.2, HIDDEN), dtype=np.float32), + ] + if not tied: + weights.append(random_weight(rng, (VOCAB, HIDDEN))) + for _ in range(N_LAYERS): + weights.extend( + [ + random_weight(rng, (HIDDEN, INTERMEDIATE)), # down + random_weight(rng, (INTERMEDIATE, HIDDEN)), # gate + np.ascontiguousarray(rng.uniform(0.8, 1.2, HIDDEN), dtype=np.float32), + np.ascontiguousarray(rng.uniform(0.8, 1.2, HIDDEN), dtype=np.float32), + random_weight(rng, (INTERMEDIATE, HIDDEN)), # up + random_weight(rng, (N_KV * HEAD_DIM, HIDDEN)), # wk + random_weight(rng, (HIDDEN, N_Q * HEAD_DIM)), # wo + random_weight(rng, (N_Q * HEAD_DIM, HIDDEN)), # wq + random_weight(rng, (N_KV * HEAD_DIM, HIDDEN)), # wv + ] + ) + return weights + + +def emit_and_compile(*, tied: bool) -> tuple[object, object]: + tag = "tied" if tied else "untied" + work = Path(tempfile.mkdtemp(prefix=f"prefill_embeddings_{tag}_")) + config_path = work / "config.json" + config_path.write_text(json.dumps(emitter_config(tied=tied)), encoding="utf-8") + + print(f"[emit] {tag}: token + embeddings prefill StableHLO", flush=True) + subprocess.run( + [ + CARGO, + "test", + "-p", + "mlxcel-xla", + "--lib", + "emitter::tests::dump_prefill_embeddings_parity_graphs", + "--", + "--ignored", + "--nocapture", + ], + cwd=REPO_ROOT, + env={ + **os.environ, + "MLXCEL_DUMP_CONFIG": str(config_path), + "MLXCEL_DUMP_DIR": str(work), + }, + check=True, + ) + + token_mlir = work / "prefill_logits.mlir" + embeddings_mlir = work / "prefill_embeddings_logits.mlir" + token_vmfb = work / "prefill_logits.vmfb" + embeddings_vmfb = work / "prefill_embeddings_logits.vmfb" + for label, source, output in [ + ("token", token_mlir, token_vmfb), + ("embeddings", embeddings_mlir, embeddings_vmfb), + ]: + print(f"[compile] {tag}/{label}: llvm-cpu", flush=True) + compile_file( + str(source), + output_file=str(output), + input_type="stablehlo", + target_backends=["llvm-cpu"], + ) + + return ( + load_vm_flatbuffer_file(str(token_vmfb), driver="local-task"), + load_vm_flatbuffer_file(str(embeddings_vmfb), driver="local-task"), + ) + + +def causal_attention_bias() -> np.ndarray: + query = np.arange(PREFILL_LP)[:, None] + key = np.arange(PREFILL_LP)[None, :] + return np.ascontiguousarray(np.where(key <= query, 0.0, MASKED), dtype=np.float32) + + +def to_host(value: object) -> np.ndarray: + host = value.to_host() if hasattr(value, "to_host") else value + return np.asarray(host, dtype=np.float32) + + +def run_case( + *, + label: str, + real_len: int, + padding_token: int, + weights: list[np.ndarray], + token_module: object, + embeddings_module: object, +) -> bool: + rng = np.random.default_rng(real_len + padding_token) + tokens = np.full(PREFILL_LP, padding_token, dtype=np.int32) + tokens[:real_len] = rng.integers(0, VOCAB, real_len, dtype=np.int32) + embeddings = np.ascontiguousarray(weights[0][tokens], dtype=np.float32) + positions = np.arange(PREFILL_LP, dtype=np.int32) + scalar_len = np.asarray(real_len, dtype=np.int32) + bias = causal_attention_bias() + + print( + f"[run] {label}: real_len={real_len} padding_token={padding_token}", + flush=True, + ) + token_out = token_module.main(*weights, tokens, positions, scalar_len) + embeddings_out = embeddings_module.main( + *weights, embeddings, positions, scalar_len, bias + ) + + names = ["logits", "kcache", "vcache"] + ok = True + for name, token_value, embeddings_value in zip( + names, token_out, embeddings_out, strict=True + ): + lhs = to_host(token_value) + rhs = to_host(embeddings_value) + max_diff = float(np.max(np.abs(lhs - rhs))) + equal = np.allclose(lhs, rhs, rtol=0.0, atol=ATOL) + ok = ok and equal + print( + f"[compare] {label}/{name}: shape={lhs.shape} " + f"max|diff|={max_diff:.3e} -> {'PASS' if equal else 'FAIL'}", + flush=True, + ) + return ok + + +def run_checkpoint(*, tied: bool, scenarios: list[tuple[str, int, int]]) -> bool: + token_module, embeddings_module = emit_and_compile(tied=tied) + weights = checkpoint_weights(tied=tied) + return all( + run_case( + label=label, + real_len=real_len, + padding_token=padding_token, + weights=weights, + token_module=token_module, + embeddings_module=embeddings_module, + ) + for label, real_len, padding_token in scenarios + ) + + +def main() -> int: + tied_ok = run_checkpoint( + tied=True, + scenarios=[ + ("tied/real_len_1", 1, 7), + ("tied/near_capacity", PREFILL_LP - 1, 9), + ], + ) + untied_ok = run_checkpoint( + tied=False, + scenarios=[("untied/nonzero_padding", 23, 5)], + ) + ok = tied_ok and untied_ok + print(f"RESULT: {'PASS' if ok else 'FAIL'} (ATOL={ATOL:g}, local-task)", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/lib/mlxcel-xla/README.md b/src/lib/mlxcel-xla/README.md index b6142050..13ff2b25 100644 --- a/src/lib/mlxcel-xla/README.md +++ b/src/lib/mlxcel-xla/README.md @@ -152,6 +152,39 @@ not fuse the in-graph dequant into the matmul on CUDA, so it regresses throughpu (see "int8 packed dequant on GB10 / CUDA" below). This is a transferable, correctness-first lever; it does not close the gap to MLX (see the perf note above). +### Prefill from embeddings + +The emitter exposes a second, non-overloaded StableHLO module named +`prefill_embeddings.main`. It begins at the shared transformer stack and returns +the same logits/argmax and per-layer K/V tensors as `prefill.main`; the existing +token module remains byte-identical at the compatibility capacity. `Lp` is the +configured static context capacity for both modules, not a separate embeddings +limit. Runtime and C-ABI wiring is intentionally a separate integration step. + +After the model weights, the embeddings entry has one deterministic argument +order: `embeddings`, `positions`, `real_len`, `attention_bias`. + +- `embeddings` is `[Lp, hidden_size]` f32 and contains post-token-embedding-scale hidden states. The graph performs no token lookup and applies neither Gemma's `sqrt(hidden_size)` scale nor an architecture embedding multiplier again. +- `positions` is `[Lp]` i32 and retains the existing one-dimensional text position contract. M-RoPE uses a later, distinct schema rather than changing this argument's meaning. +- `real_len` is an i32 scalar in `1..=Lp` and selects the output row at `real_len - 1`. +- `attention_bias` is `[Lp, Lp]` f32. `0.0` means allowed and the finite value `-1e30` means masked; NaN, infinity, and intermediate additive values are rejected. The caller supplies the complete static bucket, including padding rows. Token-parity padding retains the ordinary causal pattern on padded rows while real rows mask future and padded keys. + +`PrefillEmbeddingsInputMetadata` and the additive-bias validator provide the +pre-compilation/invocation checks for hidden size, sequence length, dtypes, +position/mask shapes, `real_len`, and mask values. The embedding table remains a +weight argument because tied checkpoints use it as the LM head; untied +checkpoints continue to consume their dedicated `lm_head` in the same weight +order. + +The pure-Rust structural/golden checks run with `cargo test -p mlxcel-xla --lib`. +The real execution fixture gathers embeddings from a deterministic tied and +untied checkpoint, compiles both modules for IREE llvm-cpu, and compares logits +plus every K/V element across one-token, nonzero-padding, and near-capacity cases: + +```bash +spike/openxla/.venv/bin/python spike/openxla/prefill_embeddings_check.py +``` + ### Scope / limits - The bundled graphs are authored for **Llama-3.2-1B-Instruct** specifically; diff --git a/src/lib/mlxcel-xla/assets/qwen2-moe-tiny/prefill_embeddings_logits.mlir b/src/lib/mlxcel-xla/assets/qwen2-moe-tiny/prefill_embeddings_logits.mlir new file mode 100644 index 00000000..a9068c3c --- /dev/null +++ b/src/lib/mlxcel-xla/assets/qwen2-moe-tiny/prefill_embeddings_logits.mlir @@ -0,0 +1,380 @@ +module @prefill_embeddings { + func.func public @main(%arg0: tensor<10x8xf32> loc("params['embed']"), %arg1: tensor<8xf32> loc("params['final_norm']"), %arg2: tensor<10x8xf32> loc("params['lm_head']"), %arg3: tensor<8xf32> loc("params['layers'][0]['in_ln']"), %arg4: tensor<8xf32> loc("params['layers'][0]['post_ln']"), %arg5: tensor<4x8xf32> loc("params['layers'][0]['wk']"), %arg6: tensor<8x8xf32> loc("params['layers'][0]['wo']"), %arg7: tensor<8x8xf32> loc("params['layers'][0]['wq']"), %arg8: tensor<4x8xf32> loc("params['layers'][0]['wv']"), %arg9: tensor<4xf32> loc("params['layers'][0]['bk']"), %arg10: tensor<8xf32> loc("params['layers'][0]['bq']"), %arg11: tensor<4xf32> loc("params['layers'][0]['bv']"), %arg12: tensor<4x8xf32> loc("params['layers'][0]['moe_router']"), %arg13: tensor<4x12x8xf32> loc("params['layers'][0]['moe_gate']"), %arg14: tensor<4x12x8xf32> loc("params['layers'][0]['moe_up']"), %arg15: tensor<4x8x12xf32> loc("params['layers'][0]['moe_down']"), %arg16: tensor<10x8xf32> loc("params['layers'][0]['moe_shared_gate']"), %arg17: tensor<10x8xf32> loc("params['layers'][0]['moe_shared_up']"), %arg18: tensor<8x10xf32> loc("params['layers'][0]['moe_shared_down']"), %arg19: tensor<1x8xf32> loc("params['layers'][0]['moe_shared_expert_gate']"), %arg20: tensor<8xf32> loc("params['layers'][1]['in_ln']"), %arg21: tensor<8xf32> loc("params['layers'][1]['post_ln']"), %arg22: tensor<4x8xf32> loc("params['layers'][1]['wk']"), %arg23: tensor<8x8xf32> loc("params['layers'][1]['wo']"), %arg24: tensor<8x8xf32> loc("params['layers'][1]['wq']"), %arg25: tensor<4x8xf32> loc("params['layers'][1]['wv']"), %arg26: tensor<4xf32> loc("params['layers'][1]['bk']"), %arg27: tensor<8xf32> loc("params['layers'][1]['bq']"), %arg28: tensor<4xf32> loc("params['layers'][1]['bv']"), %arg29: tensor<4x8xf32> loc("params['layers'][1]['moe_router']"), %arg30: tensor<4x12x8xf32> loc("params['layers'][1]['moe_gate']"), %arg31: tensor<4x12x8xf32> loc("params['layers'][1]['moe_up']"), %arg32: tensor<4x8x12xf32> loc("params['layers'][1]['moe_down']"), %arg33: tensor<10x8xf32> loc("params['layers'][1]['moe_shared_gate']"), %arg34: tensor<10x8xf32> loc("params['layers'][1]['moe_shared_up']"), %arg35: tensor<8x10xf32> loc("params['layers'][1]['moe_shared_down']"), %arg36: tensor<1x8xf32> loc("params['layers'][1]['moe_shared_expert_gate']"), %arg37: tensor<256x8xf32> loc("embeddings"), %arg38: tensor<256xi32> loc("positions"), %arg39: tensor loc("real_len"), %arg40: tensor<256x256xf32> loc("attention_bias")) -> (tensor<10xf32>, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32>) { + %0 = stablehlo.constant dense<"0x0000803F0000803F0000803F0000803F40510A3FB9FC7F3F40510A3FB9FC7F3F3311D5BEE5F27F3F3311D5BEE5F27F3F26707DBF83E27F3F26707DBF83E27F3F305527BF94CB7F3F305527BF94CB7F3F2C3C913E19AE7F3F2C3C913E19AE7F3FB8CD753F128A7F3FB8CD753F128A7F3FBDFF403F815F7F3FBDFF403F815F7F3FF6FD14BE662E7F3FF6FD14BE662E7F3FD53F69BFC2F67E3FD53F69BFC2F67E3F64CD56BF98B87E3F64CD56BF98B87E3F7205913BE8737E3F7205913BE8737E3FD006583FB5287E3FD006583FB5287E3F6F4E683F00D77D3F6F4E683F00D77D3FD7040C3ECC7E7D3FD7040C3ECC7E7D3FE87A42BF1A207D3FE87A42BF1A207D3F2C2975BFEDBA7C3F2C2975BFEDBA7C3F36E28CBE494F7C3F36E28CBE494F7C3F840A293F2EDD7B3F840A293F2EDD7B3FBF1B7D3FA1647B3FBF1B7D3FA1647B3F22F0D03EA5E57A3F22F0D03EA5E57A3FFC370CBF3C607A3FFC370CBF3C607A3F6FFD7FBF6AD4793F6FFD7FBF6AD4793FBF6708BF3342793FBF6708BF3342793FFE2DD93E9BA9783FFE2DD93E9BA9783F78BF7D3FA50A783F78BF7D3FA50A783F819C253F5565773F819C253F5565773F389395BEB1B9763F389395BEB1B9763F576D76BFBB07763F576D76BFBB07763FB3803FBF784F753FB3803FBF784F753F18F41D3EEF90743F18F41D3EEF90743F8E2C6A3F22CC733F8E2C6A3F22CC733FAA8F553F1801733FAA8F553F1801733FB78659BCD52F723FB78659BCD52F723FE73B59BF5F58713FE73B59BF5F58713F5F5867BFBB7A703F5F5867BFBB7A703FEA0803BEEF966F3FEA0803BEEF966F3F2DF2433F01AD6E3F2DF2433F01AD6E3FB57F743FF7BC6D3FB57F743FF7BC6D3F6C85883ED7C66C3F6C85883ED7C66C3F74BC2ABFA7CA6B3F74BC2ABFA7CA6B3F44C27CBF6EC86A3F44C27CBF6EC86A3FE0CACCBE32C0693FE0CACCBE32C0693FE81B0E3FFBB1683FE81B0E3FFBB1683FBBF57F3FCE9D673FBBF57F3FCE9D673F807B063FB483663F807B063FB483663F6D46DDBEB263653F6D46DDBEB263653FB3097EBFD23D643FB3097EBFD23D643F80E023BF1912633F80E023BF1912633F44E7993E91E0613F44E7993E91E0613F0308773F40A9603F0308773F40A9603FD1FD3D3F2F6C5F3FD1FD3D3F2F6C5F3F0EE726BE66295E3F0EE726BE66295E3F95146BBFEDE05C3F95146BBFEDE05C3FA64D54BFCC925B3FA64D54BFCC925B3F2C43B53C0C3F5A3F2C43B53C0C3F5A3FA26C5A3FB6E5583FA26C5A3FB6E5583FAC5D663FD386573FAC5D663FD386573FB714F43D6B22563FB714F43D6B22563F836545BF87B8543F836545BF87B8543F56D173BF3249533F56D173BF3249533FE62584BE74D4513FE62584BE74D4513FF76A2C3F575A503FF76A2C3F575A503FB8637C3FE4DA4E3FB8637C3FE4DA4E3F83A1C83E26564D3F83A1C83E26564D3FFAFC0FBF26CC4B3FFAFC0FBF26CC4B3FE5E87FBFEE3C4A3FE5E87FBFEE3C4A3F908C04BF89A8483F908C04BF89A8483F6D5AE13E010F473F6D5AE13E010F473FD54E7E3F6170453FD54E7E3F6170453F3521223FB3CC433F3521223FB3CC433F3A389EBE0124423F3A389EBE0124423FBC9D77BF5776403FBC9D77BF5776403F20773CBFC0C33E3F20773CBFC0C33E3FACD62F3E470C3D3FACD62F3E470C3D3FE4F76B3FF64F3B3FE4F76B3FF64F3B3F6107533FDA8E393F6107533FDA8E393F5ABFFDBCFEC8373F5ABFFDBCFEC8373FFC985BBF6EFE353FFC985BBF6EFE353F5A5E65BF352F343F5A5E65BF352F343FB512E2BD5F5B323FB512E2BD5F5B323FE4D4463FF882303FE4D4463FF882303F141E733F0DA62E3F141E733F0DA62E3F72877F3EA8C42C3F72877F3EA8C42C3F05162EBFD8DE2A3F05162EBFD8DE2A3F1C007CBFA8F4283F1C007CBFA8F4283F1F74C4BE2406273F1F74C4BE2406273F28DB113F5A13253F28DB113F5A13253FEDD67F3F561C233FEDD67F3F561C233FF69A023F2621213FF69A023F2621213FE869E5BED5211F3FE869E5BED5211F3FDE8E7EBF711E1D3FDE8E7EBF711E1D3FAA5E20BF08171B3FAA5E20BF08171B3F0486A23EA60B193F0486A23EA60B193F7D2E783F59FC163F7D2E783F59FC163FA8EC3A3F2FE9143FA8EC3A3F2FE9143FC2C238BE35D2123FC2C238BE35D2123F78D66CBF78B7103F78D66CBF78B7103FE0BC51BF08990E3FE0BC51BF08990E3F381B233DF0760C3F381B233DF0760C3FEEC05C3F40510A3FEEC05C3F40510A3F6E5A643F0628083F6E5A643F0628083F2A0CD03D50FB053F2A0CD03D50FB053F494048BF2BCB033F494048BF2BCB033FF16572BFA797013FF16572BFA797013FF8BD76BEA2C1FE3EF8BD76BEA2C1FE3E95BD2F3F714DFA3E95BD2F3F714DFA3E71977B3FD8D2F53E71977B3FD8D2F53ECB42C03EF451F13ECB42C03EF451F13E6AB613BFE2CAEC3E6AB613BFE2CAEC3ED4BF7FBFC03DE83ED4BF7FBFC03DE83EBFA600BFADAAE33EBFA600BFADAAE33EC874E93EC511DF3EC874E93EC511DF3ECBC97E3F2873DA3ECBC97E3F2873DA3EE7981E3FF2CED53EE7981E3FF2CED53E8BD0A6BE4425D13E8BD0A6BE4425D13E44BA78BF3B76CC3E44BA78BF3B76CC3E6F5E39BFF6C1C73E6F5E39BFF6C1C73E25AB413E9408C33E25AB413E9408C33E4CB06D3F344ABE3E4CB06D3F344ABE3E2A6E503FF586B93E2A6E503FF586B93E7E5347BDF5BEB43E7E5347BDF5BEB43E73E45DBF56F2AF3E73E45DBF56F2AF3EEE5163BF3521AB3EEE5163BF3521AB3E7301BEBDB34BA63E7301BEBDB34BA63EA9A7493FEF71A13EA9A7493FEF71A13EF2A8713F09949C3EF2A8713F09949C3E8BEF6D3E20B2973E8BEF6D3E20B2973E9F6131BF56CC923E9F6131BF56CC923EBB297BBFCAE28D3EBB297BBFCAE28D3E9C0DBCBE9BF5883E9C0DBCBE9BF5883EB68E153FEB04843EB68E153FEB04843E99A37F3FB5217E3E99A37F3FB5217E3EE45FFD3E1133743EE45FFD3E1133743EFA7AEDBE2D3E6A3EFA7AEDBE2D3E6A3E9DFF7EBF4A43603E9DFF7EBF4A43603EF6CF1CBFA942563EF6CF1CBFA942563EBA17AB3E8C3C4C3EBA17AB3E8C3C4C3E0E41793F3531423E0E41793F3531423E7FCC373FE520383E7FCC373FE520383EA48F4ABEDE0B2E3EA48F4ABEDE0B2E3E5B856EBF62F2233E5B856EBF62F2233E461B4FBFB5D4193E461B4FBFB5D4193EC5876B3D17B30F3EC5876B3D17B30F3E84035F3FCB8D053E84035F3FCB8D053EDF44623F28CAF63DDF44623F28CAF63DEDF2AB3D6972E23DEDF2AB3D6972E23DFD0A4BBFDD14CE3DFD0A4BBFDD14CE3D1AE770BF0BB2B93D1AE770BF0BB2B93D591C65BE794AA53D591C65BE794AA53D1A02333FAADE903D1A02333FAADE903DFCB67A3F4EDE783DFCB67A3F4EDE783DA7D4B73EE8F84F3DA7D4B73EE8F84F3D016417BF2F0E273D016417BF2F0E273D3E827FBF5F3EFC3C3E827FBF5F3EFC3C366DF9BEEA59AA3C366DF9BEEA59AA3C697CF13E32E2303C697CF13E32E2303C51307F3F95C0503A51307F3F95C0503AE0031B3F75CA16BCE0031B3F75CA16BC7A5BAFBE8C4E9DBC7A5BAFBE8C4E9DBCD8C279BFD633EFBCD8C279BFD633EFBCDF3636BF808920BDDF3636BF808920BD1470533EFA7449BD1470533EFA7449BDA2556F3F4B5B72BDA2556F3F4B5B72BD3BC44D3FB49D8DBD3BC44D3FB49D8DBDA9DB87BD220AA2BDA9DB87BD220AA2BD1C1E60BF6A72B6BD1C1E60BF6A72B6BD463361BF07D6CABD463361BF07D6CABDF3E099BD7234DFBDF3E099BD7234DFBD3F6A4C3F278DF3BD3F6A4C3F278DF3BD6D20703FD0EF03BE6D20703FD0EF03BE8E445C3EAB150EBE8E445C3EAB150EBEFE9E34BFE43718BEFE9E34BFE43718BE353F7ABF365622BE353F7ABF365622BE0298B3BE61702CBE0298B3BE61702CBE4336193F228636BE4336193F228636BEC35B7F3F379740BEC35B7F3F379740BE8875F53E5EA34ABE8875F53E5EA34ABEFF78F5BE54AA54BEFF78F5BE54AA54BEE75B7FBFD9AB5EBEE75B7FBFD9AB5EBEAE3419BFAAA768BEAE3419BFAAA768BEB69BB33E879D72BEB69BB33E879D72BE9F3F7A3F2E8D7CBE9F3F7A3F2E8D7CBE979D343F2F3B83BE979D343F2F3B83BE464C5CBE6B2C88BE464C5CBE6B2C88BE1C2170BF2A1A8DBE1C2170BF2A1A8DBE0F694CBF4D0492BE0F694CBF4D0492BEB6F0993DB3EA96BEB6F0993DB3EA96BE3634613F3BCD9BBE3634613F3BCD9BBE281D603FC7ABA0BE281D603FC7ABA0BEE4CB873D3686A5BEE4CB873D3686A5BE67C54DBF685CAABE67C54DBF685CAABEEF546FBF3D2EAFBEEF546FBF3D2EAFBE586853BE97FBB3BE586853BE97FBB3BE4238363F55C4B8BE4238363F55C4B8BE69C2793F5888BDBE69C2793F5888BDBEC357AF3E8047C2BEC357AF3E8047C2BE72051BBFB001C7BE72051BBFB001C7BE29307FBFC8B6CBBE29307FBFC8B6CBBEED78F1BEA866D0BEED78F1BEA866D0BEAA70F93E3311D5BEAA70F93E3311D5BE5D827F3F49B6D9BE5D827F3F49B6D9BE6962173FCD55DEBE6962173FCD55DEBE57D8B7BE9FEFE2BE57D8B7BE9FEFE2BE62B77ABFA283E7BE62B77ABFA283E7BEB00033BFB811ECBEB00033BFB811ECBE0D24653EC399F0BE0D24653EC399F0BEC5E7703FA51BF5BEC5E7703FA51BF5BEC9094B3F4197F9BEC9094B3F4197F9BEAC02ACBD790CFEBEAC02ACBD790CFEBECB4562BF983D01BFCB4562BF983D01BF8C025FBFA57103BF8C025FBFA57103BF36686BBD54A205BF36686BBD54A205BF6F1C4F3F97CF07BF6F1C4F3F97CF07BFA4846E3F61F909BFA4846E3F61F909BFE5874A3EA21F0CBFE5874A3EA21F0CBFDFCD37BF4D420EBFDFCD37BF4D420EBF9A4079BF536110BF9A4079BF536110BF0014ABBEA77C12BF0014ABBEA77C12BF86D11C3F3C9414BF86D11C3F3C9414BF70FF7E3F02A816BF70FF7E3F02A816BF7A77ED3EEEB718BF7A77ED3EEEB718BF5363FDBEF0C31ABF5363FDBEF0C31ABFB4A37FBFFCCB1CBFB4A37FBFFCCB1CBF1B8D15BF05D01EBF1B8D15BF05D01EBF4911BC3EFDCF20BF4911BC3EFDCF20BF1D2A7B3FD6CB22BF1D2A7B3FD6CB22BF3260313F85C324BF3260313F85C324BF3BF76DBEFDB626BF3BF76DBEFDB626BF99A971BF2FA628BF99A971BF2FA628BF71A649BF11912ABF71A649BF11912ABF3011BE3D94772CBF3011BE3D94772CBFD752633FAE592EBFD752633FAE592EBF76E35D3F503730BF76E35D3F503730BFEC33473D701032BFEC33473D701032BF506F50BF01E533BF506F50BF01E533BF90AF6DBFF7B435BF90AF6DBFF7B435BF62A341BE468037BF62A341BE468037BFCC5F393FE34639BFCC5F393FE34639BFCCB9783FC1083BBFCCB9783FC1083BBFCFCCA63ED6C53CBFCFCCA63ED6C53CBF749A1EBF157E3EBF749A1EBF157E3EBF9AC97EBF743140BF9AC97EBF743140BF4471E9BEE8DF41BF4471E9BEE8DF41BF74A8003F658943BF74A8003F658943BFEABF7F3FE02D45BFEABF7F3FE02D45BFCDB4133F4FCD46BFCDB4133F4FCD46BF7546C0BEA86748BF7546C0BEA86748BFCF977BBFDFFC49BFCF977BBFDFFC49BF25BC2FBFEA8C4BBF25BC2FBFEA8C4BBFA3C5763EBF174DBFA3C5763EBF174DBF9466723F559D4EBF9466723F559D4EBF0E3F483FA01D50BF0E3F483FA01D50BFE31BD0BD979851BFE31BD0BD979851BF535B64BF310E53BF535B64BF310E53BFEEBF5CBF637E54BFEEBF5CBF637E54BF"> : tensor<256x4xf32> + %1 = stablehlo.constant dense<"0x00000000000000000000000000000000A46A573F57D6233CA46A573F57D6233CB7C7683F3ED4A33CB7C7683F3ED4A33CC381103E20B9F53CC381103E20B9F53CCFBD41BFDBCB233DCFBD41BFDBCB233D107C75BFF5B64C3D107C75BFF5B64C3D8C0F8FBED19C753D8C0F8FBED19C753D4630283F323E8F3D4630283F323E8F3D95467D3F51AAA33D95467D3F51AAA33D3201D33E3F12B83D3201D33E3F12B83DF8440BBF7675CC3DF8440BBF7675CC3D5CFF7FBF72D3E03D5CFF7FBF72D3E03DD85C09BFAC2BF53DD85C09BFAC2BF53D2220D73ED0BE043E2220D73ED0BE043E72987D3F64E40E3E72987D3F64E40E3E4479263F4F06193E4479263F4F06193E106893BE5024233E106893BE5024233E251E76BF233E2D3E251E76BF233E2D3EB34040BF8753373EB34040BF8753373E6979193E3A64413E6979193E3A64413EC8B6693FF96F4B3EC8B6693FF96F4B3E102F563F8376553E102F563F8376553E150511BC96775F3E150511BC96775F3EE7A158BFF172693EE7A158BFF172693EFCD367BF5168733EFCD367BF5168733E388707BE77577D3E388707BE77577D3E0837433F10A0833E0837433F10A0833E0ED5743F0691883E0ED5743F0691883E2AB48A3E7D7E8D3E2AB48A3E7D7E8D3EE9E329BF5468923EE9E329BF5468923EA4EF7CBF6D4E973EA4EF7CBF6D4E973E06DECEBEA5309C3E06DECEBEA5309C3E4D2A0D3FDE0EA13E4D2A0D3FDE0EA13E39FA7F3FF8E8A53E39FA7F3FF8E8A53EF771073FD2BEAA3EF771073FD2BEAA3EC23ADBBE4D90AF3EC23ADBBE4D90AF3E38E57DBF4A5DB43E38E57DBF4A5DB43EEABE24BFA825B93EEABE24BFA825B93E9FBD973E4AE9BD3E9FBD973E4AE9BD3E4BBB763F0EA8C23E4BBB763F0EA8C23EBCBF3E3FD761C73EBCBF3E3FD761C73EFB6D22BE8616CC3EFB6D22BE8616CC3E28A16ABFFAC5D03E28A16ABFFAC5D03E31EF54BF1770D53E31EF54BF1770D53EA103913CBD14DA3EA103913CBD14DA3ED0D4593FCDB3DE3ED0D4593FCDB3DE3E9ADB663F2B4DE33E9ADB663F2B4DE33EE813FD3DB6E0E73EE813FD3DB6E0E73E56AC44BF526EEC3E56AC44BF526EEC3E222974BFE0F5F03E222974BFE0F5F03EFF5586BE4477F53EFF5586BE4477F53E24942B3F5EF2F93E24942B3F5EF2F93EA0937C3F1367FE3EA0937C3F1367FE3EB4B6CA3EA26A013FB4B6CA3EA26A013FCD0C0FBF6B9E033FCD0C0FBF6B9E033FF4EF7FBFD5CE053FF4EF7FBFD5CE053F5E8405BFD2FB073F5E8405BFD2FB073FFD50DF3E53250A3FFD50DF3E53250A3FE72C7E3F4C4B0C3FE72C7E3F4C4B0C3F4301233FAD6D0E3F4301233FAD6D0E3F23109CBE698C103F23109CBE698C103F7F5377BF71A7123F7F5377BF71A7123FF23A3DBFB8BE143FF23A3DBFB8BE143F4B5F2B3E31D2163F4B5F2B3E31D2163FD4866B3FCDE1183FD4866B3FCDE1183F0BAB533F7FED1A3F0BAB533F7FED1A3FCE81D9BC3AF51C3FCE81D9BC3AF51C3F5B035BBFF0F81E3F5B035BBFF0F81E3F97DE65BF94F8203F97DE65BF94F8203F4D14EBBD1AF4223F4D14EBBD1AF4223FB31D463F73EB243FB31D463F73EB243F5178733F94DE263F5178733F94DE263F23F5813E6FCD283F23F5813E6FCD283FED402DBFF8B72A3FED402DBFF8B72A3F8C327CBF229E2C3F8C327CBF229E2C3F518BC6BEE17F2E3F518BC6BEE17F2E3F6EEC103F285D303F6EEC103F285D303F8EE07F3FEB35323F8EE07F3FEB35323F1894033F1F0A343F1894033F1F0A343FBC62E3BEB6D9353FBC62E3BEB6D9353F7D6F7EBFA6A4373F7D6F7EBFA6A4373F574021BFE26A393F574021BFE26A393F865FA03E602C3B3F865FA03E602C3B3FBCE6773F12E93C3FBCE6773F12E93C3F5CB23B3FEEA03E3F5CB23B3FEEA03E3F2B4D34BEEA53403F2B4D34BEEA53403FC6676CBFF801423FC6676CBFF801423FA76252BF10AB433FA76252BF10AB433FD0FD103D254F453FD0FD103D254F453F822D5C3F2DEE463F822D5C3F2DEE463FF7DC643F1D88483FF7DC643F1D88483FFB0FD93DEB1C4A3FFB0FD93DEB1C4A3F178B47BF8CAC4B3F178B47BF8CAC4B3F9EC272BFF7364D3F9EC272BFF7364D3F56237BBE21BC4E3F56237BBE21BC4E3F3DEA2E3F003C503F3DEA2E3F003C503F68CC7B3F8AB6513F68CC7B3F8AB6513FF25BC23EB62B533FF25BC23EB62B533F28C912BF7A9B543F28C912BF7A9B543F05CC7FBFCC05563F05CC7FBFCC05563F2EA101BFA46A573F2EA101BFA46A573FED6FE73EF9C9583FED6FE73EF9C9583FF8AC7E3FC0235A3FF8AC7E3FC0235A3F2F7C1F3FF2775B3F2F7C1F3FF2775B3FB1ABA4BE86C65C3FB1ABA4BE86C65C3F007578BF730F5E3F007578BF730F5E3F03263ABFB0525F3F03263ABFB0525F3F6D373D3E3690603F6D373D3E3690603FFA436D3FFDC7613FFA436D3FFDC7613F0B16513FFBF9623F0B16513FFBF9623FCF3735BD2A26643FCF3735BD2A26643F3E535DBF824C653F3E535DBF824C653FC0D663BFFC6C663FC0D663BFFC6C663F4F07C7BD8F87673F4F07C7BD8F87673F7AF4483F349C683F7AF4483F349C683F0D08723FE6AA693F0D08723FE6AA693F5D57723E9CB36A3F5D57723E9CB36A3F0B9030BF50B66B3F0B9030BF50B66B3F38617BBFFBB26C3F38617BBFFBB26C3FAE28BEBE97A96D3FAE28BEBE97A96D3FEFA2143F1D9A6E3FEFA2143F1D9A6E3F5BB27F3F88846F3F5BB27F3F88846F3F5557FF3ED168703F5557FF3ED168703F7878EBBEF246713F7878EBBEF246713F58E57EBFE61E723F58E57EBFE61E723FD4B41DBFA8F0723FD4B41DBFA8F0723F8FF4A83E31BC733F8FF4A83E31BC733F49FE783F7D81743F49FE783F7D81743FED95383F8740753FED95383F8740753FE41D46BE49F9753FE41D46BE49F9753F6C1B6EBFC0AB763F6C1B6EBFC0AB763F3DC54FBFE657773F3DC54FBFE657773F2D6E593DB6FD773F2D6E593DB6FD773F8A745E3F2E9D783F8A745E3F2E9D783FF8CB623F4836793FF8CB623F4836793FA4FAB43D02C9793FA4FAB43D02C9793FD5594ABF56557A3FD5594ABF56557A3FA14871BF41DB7A3FA14871BF41DB7A3F888669BEC15A7B3F888669BEC15A7B3F4F32323FD1D37B3F4F32323FD1D37B3FFCF07A3F6F467C3FFCF07A3F6F467C3F99F1B93E98B27C3F99F1B93E98B27C3FBC7916BF48187D3FBC7916BF48187D3F90937FBF7E777D3F90937FBF7E777D3F2F67FBBE37D07D3F2F67FBBE37D07D3F4B7CEF3E70227E3F4B7CEF3E70227E3F9B187F3F286E7E3F9B187F3F286E7E3F4FEA1B3F5DB37E3F4FEA1B3F5DB37E3F093AADBE0CF27E3F093AADBE0CF27E3F938279BF352A7F3F938279BF352A7F3F240237BFD55B7F3F240237BFD55B7F3F61004F3EEC867F3F61004F3EEC867F3F18EE6E3F78AB7F3F18EE6E3F78AB7F3F45704E3F78C97F3F45704E3F78C97F3F2EA07DBDEDE07F3F2EA07DBDEDE07F3F60915FBFD4F17F3F60915FBFD4F17F3FA3BC61BF2EFC7F3FA3BC61BF2EFC7F3F59EAA2BDFBFF7F3F59EAA2BDFBFF7F3F21BB4B3F39FD7F3F21BB4B3F39FD7F3F5D84703FEBF37F3F5D84703FEBF37F3F04B1603E0EE47F3F04B1603E0EE47F3FFFD033BFA5CD7F3FFFD033BFA5CD7F3FB97B7ABFAFB07F3FB97B7ABFAFB07F3FC9B6B5BE2E8D7F3FC9B6B5BE2E8D7F3F844D183F22637F3F844D183F22637F3FA46F7F3F8C327F3FA46F7F3F8C327F3FFE71F73E6EFB7E3FFE71F73E6EFB7E3F507BF3BEC8BD7E3F507BF3BEC8BD7E3FC0467FBF9E797E3FC0467FBF9E797E3FAA1C1ABFEF2E7E3FAA1C1ABFEF2E7E3F0A7CB13EBFDD7D3F0A7CB13EBFDD7D3FDC017A3F0F867D3FDC017A3F0F867D3FAF6A353FE1277D3FAF6A353FE1277D3FB7DE57BE39C37C3FB7DE57BE39C37C3FF9BB6FBF18587C3FF9BB6FBF18587C3F28174DBF81E67B3F28174DBF81E67B3F8CE6903D776E7B3F8CE6903D776E7B3FB9A9603FFEEF7A3FB9A9603FFEEF7A3FC7A8603F186B7A3FC7A8603F186B7A3FC9D6903DC9DF793FC9D6903DC9DF793F57184DBF144E793F57184DBF144E793F47BB6FBFFEB5783F47BB6FBFFEB5783FFED657BE8917783FFED657BE8917783F146C353FBB72773F146C353FBB72773F6F017A3F97C7763F6F017A3F97C7763F5578B13E2216763F5578B13E2216763F3E1E1ABF605E753F3E1E1ABF605E753F9A467FBF56A0743F9A467FBF56A0743FD777F3BE09DC733FD777F3BE09DC733F7375F73E7D11733F7375F73E7D11733FC66F7F3FB940723FC66F7F3FB940723FED4B183FC169713FED4B183FC169713F7BBAB5BE9B8C703F7BBAB5BE9B8C703F217C7ABF4DA96F3F217C7ABF4DA96F3F97CF33BFDCBF6E3F97CF33BFDCBF6E3FB9B8603E4ED06D3FB9B8603E4ED06D3F0B85703FAADA6C3F0B85703FAADA6C3FEFB94B3FF5DE6B3FEFB94B3FF5DE6B3F1AFAA2BD37DD6A3F1AFAA2BD37DD6A3F92BD61BF75D5693F92BD61BF75D5693F69905FBFB7C7683F69905FBFB7C7683FA2807DBD04B4673FA2807DBD04B4673F70714E3F629A663F70714E3F629A663F62ED6E3FD97A653F62ED6E3FD97A653FA4F84E3E7055643FA4F84E3E7055643F860337BF2E2A633F860337BF2E2A633F228279BF1CF9613F228279BF1CF9613F5136ADBE41C2603F5136ADBE41C2603FE0EB1B3FA5855F3FE0EB1B3FA5855F3F70187F3F50435E3F70187F3F50435E3FCD78EF3E4AFB5C3FCD78EF3E4AFB5C3FA06AFBBE9CAD5B3FA06AFBBE9CAD5B3FAD937FBF4F5A5A3FAD937FBF4F5A5A3F237816BF6A01593F237816BF6A01593F47F5B93EF8A2573F47F5B93EF8A2573F61F17A3F003F563F61F17A3F003F563FE430323F8CD5543FE430323F8CD5543F3A8E69BEA566533F3A8E69BEA566533F4A4971BF55F2513F4A4971BF55F2513F9F584ABFA578503F9F584ABFA578503F620AB53D9FF94E3F620AB53D9FF94E3FE3CC623F4C754D3FE3CC623F4C754D3F90735E3FB7EB4B3F90735E3FB7EB4B3F9D4E593DE95C4A3F9D4E593DE95C4A3F65C64FBFEEC8483F65C64FBFEEC8483FB31A6EBFCE2F473FB31A6EBFCE2F473F231646BE9591453F231646BE9591453F4C97383F4DEE433F4C97383F4DEE433FD3FD783F0246423FD3FD783F0246423FD4F0A83EBD98403FD4F0A83EBD98403F62B61DBF89E63E3F62B61DBF89E63E3F29E57EBF732F3D3F29E57EBF732F3D3FF674EBBE85733B3FF674EBBE85733B3FC15AFF3ECBB2393FC15AFF3ECBB2393F73B27F3F4FED373F73B27F3F4FED373F53A1143F1E23363F53A1143F1E23363F592CBEBE4454343F592CBEBE4454343F97617BBFCB80323F97617BBFCB80323F9D8E30BFC1A8303F9D8E30BFC1A8303F0B5F723E31CC2E3F0B5F723E31CC2E3FB208723F28EB2C3FB208723F28EB2C3F40F3483FB1052B3F40F3483FB1052B3F0917C7BDD91B293F0917C7BDD91B293FA7D763BFAE2D273FA7D763BFAE2D273F40525DBF3A3B253F40525DBF3A3B253F3C1835BD8C44233F3C1835BD8C44233F2F17513FB049213F2F17513FB049213F3C436D3FB34A1F3F3C436D3FB34A1F3FA92F3D3EA2471D3FA92F3D3EA2471D3F5E273ABF8A401B3F5E273ABF8A401B3F867478BF7835193F867478BF7835193FF3A7A4BE7B26173FF3A7A4BE7B26173FBA7D1F3F9F13153FBA7D1F3F9F13153FC5AC7E3FF2FC123FC5AC7E3FF2FC123F666CE73E82E2103F666CE73E82E2103FE2A201BF5CC40E3FE2A201BF5CC40E3F"> : tensor<256x4xf32> + %2 = stablehlo.constant dense<0x00000000> : tensor + %3 = stablehlo.constant dense<0x3F800000> : tensor + %4 = stablehlo.constant dense<0xFF800000> : tensor + %5 = stablehlo.constant dense<0xF149F2CA> : tensor + %6 = stablehlo.constant dense<0x358637BD> : tensor + %7 = stablehlo.constant dense<0x41000000> : tensor + %8 = stablehlo.constant dense<0x3F000000> : tensor + %9 = stablehlo.constant dense<0> : tensor + %10 = stablehlo.constant dense<0> : tensor + %11 = stablehlo.constant dense<1> : tensor + %12 = stablehlo.reshape %arg38 : (tensor<256xi32>) -> tensor<256x1xi32> + %13 = "stablehlo.gather"(%0, %12) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %14 = "stablehlo.gather"(%1, %12) <{dimension_numbers = #stablehlo.gather, slice_sizes = array}> : (tensor<256x4xf32>, tensor<256x1xi32>) -> tensor<256x4xf32> + %15 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %16 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<2x256x1x4xf32> + %17 = stablehlo.multiply %arg37, %arg37 : tensor<256x8xf32> + %18 = stablehlo.reduce(%17 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %19 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %20 = stablehlo.divide %18, %19 : tensor<256xf32> + %21 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %22 = stablehlo.add %20, %21 : tensor<256xf32> + %23 = stablehlo.rsqrt %22 : tensor<256xf32> + %24 = stablehlo.broadcast_in_dim %23, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %25 = stablehlo.multiply %arg37, %24 : tensor<256x8xf32> + %26 = stablehlo.broadcast_in_dim %arg3, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %27 = stablehlo.multiply %25, %26 : tensor<256x8xf32> + %28 = stablehlo.dot_general %27, %arg7, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %29 = stablehlo.broadcast_in_dim %arg10, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %30 = stablehlo.add %28, %29 : tensor<256x8xf32> + %31 = stablehlo.reshape %30 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %32 = stablehlo.dot_general %27, %arg5, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %33 = stablehlo.broadcast_in_dim %arg9, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %34 = stablehlo.add %32, %33 : tensor<256x4xf32> + %35 = stablehlo.reshape %34 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %36 = stablehlo.dot_general %27, %arg8, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %37 = stablehlo.broadcast_in_dim %arg11, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %38 = stablehlo.add %36, %37 : tensor<256x4xf32> + %39 = stablehlo.reshape %38 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %40 = stablehlo.broadcast_in_dim %13, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %41 = stablehlo.broadcast_in_dim %14, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %42 = stablehlo.multiply %31, %40 : tensor<256x2x4xf32> + %43 = stablehlo.slice %31 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %44 = stablehlo.slice %31 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %45 = stablehlo.negate %44 : tensor<256x2x2xf32> + %46 = stablehlo.concatenate %45, %43, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %47 = stablehlo.multiply %46, %41 : tensor<256x2x4xf32> + %48 = stablehlo.add %42, %47 : tensor<256x2x4xf32> + %49 = stablehlo.broadcast_in_dim %13, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %50 = stablehlo.broadcast_in_dim %14, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %51 = stablehlo.multiply %35, %49 : tensor<256x1x4xf32> + %52 = stablehlo.slice %35 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %53 = stablehlo.slice %35 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %54 = stablehlo.negate %53 : tensor<256x1x2xf32> + %55 = stablehlo.concatenate %54, %52, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %56 = stablehlo.multiply %55, %50 : tensor<256x1x4xf32> + %57 = stablehlo.add %51, %56 : tensor<256x1x4xf32> + %58 = stablehlo.reshape %57 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %59 = stablehlo.dynamic_update_slice %15, %58, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %60 = stablehlo.reshape %39 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %61 = stablehlo.dynamic_update_slice %16, %60, %10, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %62 = stablehlo.reshape %48 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %63 = stablehlo.dot_general %62, %57, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %64 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %65 = stablehlo.multiply %63, %64 : tensor<1x256x2x256xf32> + %66 = stablehlo.broadcast_in_dim %arg40, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %67 = stablehlo.add %65, %66 : tensor<1x256x2x256xf32> + %68 = stablehlo.reduce(%67 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %69 = stablehlo.broadcast_in_dim %68, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %70 = stablehlo.subtract %67, %69 : tensor<1x256x2x256xf32> + %71 = stablehlo.exponential %70 : tensor<1x256x2x256xf32> + %72 = stablehlo.reduce(%71 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %73 = stablehlo.broadcast_in_dim %72, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %74 = stablehlo.divide %71, %73 : tensor<1x256x2x256xf32> + %75 = stablehlo.dot_general %74, %39, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %76 = stablehlo.transpose %75, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %77 = stablehlo.reshape %76 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %78 = stablehlo.dot_general %77, %arg6, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %79 = stablehlo.add %arg37, %78 : tensor<256x8xf32> + %80 = stablehlo.multiply %79, %79 : tensor<256x8xf32> + %81 = stablehlo.reduce(%80 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %82 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %83 = stablehlo.divide %81, %82 : tensor<256xf32> + %84 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %85 = stablehlo.add %83, %84 : tensor<256xf32> + %86 = stablehlo.rsqrt %85 : tensor<256xf32> + %87 = stablehlo.broadcast_in_dim %86, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %88 = stablehlo.multiply %79, %87 : tensor<256x8xf32> + %89 = stablehlo.broadcast_in_dim %arg4, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %90 = stablehlo.multiply %88, %89 : tensor<256x8xf32> + %91 = stablehlo.dot_general %90, %arg12, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %92 = stablehlo.reduce(%91 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %93 = stablehlo.broadcast_in_dim %92, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %94 = stablehlo.subtract %91, %93 : tensor<256x4xf32> + %95 = stablehlo.exponential %94 : tensor<256x4xf32> + %96 = stablehlo.reduce(%95 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %97 = stablehlo.broadcast_in_dim %96, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %98 = stablehlo.divide %95, %97 : tensor<256x4xf32> + %99 = stablehlo.iota dim = 0 : tensor<4xi32> + %100 = stablehlo.broadcast_in_dim %99, dims = [1] : (tensor<4xi32>) -> tensor<256x4xi32> + %101 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x4xf32> + %102 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor) -> tensor<256x4xf32> + %103 = stablehlo.constant dense<0> : tensor + %104 = stablehlo.constant dense<0xFF800000> : tensor + %105 = stablehlo.iota dim = 1 : tensor<256x4xi32> + %106:2 = stablehlo.reduce(%98 init: %104), (%105 init: %103) across dimensions = [1] : (tensor<256x4xf32>, tensor<256x4xi32>, tensor, tensor) -> (tensor<256xf32>, tensor<256xi32>) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %107 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %108 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %109 = stablehlo.or %107, %108 : tensor + %110 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %111 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %112 = stablehlo.and %110, %111 : tensor + %113 = stablehlo.or %109, %112 : tensor + %114 = stablehlo.select %109, %amv_l, %amv_r : tensor, tensor + %115 = stablehlo.select %113, %ami_l, %ami_r : tensor, tensor + stablehlo.return %114, %115 : tensor, tensor + } + %116 = stablehlo.broadcast_in_dim %106#1, dims = [0] : (tensor<256xi32>) -> tensor<256x4xi32> + %117 = stablehlo.compare EQ, %100, %116, SIGNED : (tensor<256x4xi32>, tensor<256x4xi32>) -> tensor<256x4xi1> + %118 = stablehlo.select %117, %98, %101 : tensor<256x4xi1>, tensor<256x4xf32> + %119 = stablehlo.reduce(%118 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %120 = stablehlo.select %117, %102, %98 : tensor<256x4xi1>, tensor<256x4xf32> + %121 = stablehlo.constant dense<0> : tensor + %122 = stablehlo.constant dense<0xFF800000> : tensor + %123 = stablehlo.iota dim = 1 : tensor<256x4xi32> + %124:2 = stablehlo.reduce(%120 init: %122), (%123 init: %121) across dimensions = [1] : (tensor<256x4xf32>, tensor<256x4xi32>, tensor, tensor) -> (tensor<256xf32>, tensor<256xi32>) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %125 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %126 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %127 = stablehlo.or %125, %126 : tensor + %128 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %129 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %130 = stablehlo.and %128, %129 : tensor + %131 = stablehlo.or %127, %130 : tensor + %132 = stablehlo.select %127, %amv_l, %amv_r : tensor, tensor + %133 = stablehlo.select %131, %ami_l, %ami_r : tensor, tensor + stablehlo.return %132, %133 : tensor, tensor + } + %134 = stablehlo.broadcast_in_dim %124#1, dims = [0] : (tensor<256xi32>) -> tensor<256x4xi32> + %135 = stablehlo.compare EQ, %100, %134, SIGNED : (tensor<256x4xi32>, tensor<256x4xi32>) -> tensor<256x4xi1> + %136 = stablehlo.select %135, %120, %101 : tensor<256x4xi1>, tensor<256x4xf32> + %137 = stablehlo.reduce(%136 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %138 = stablehlo.select %135, %102, %120 : tensor<256x4xi1>, tensor<256x4xf32> + %139 = stablehlo.add %119, %137 : tensor<256xf32> + %140 = stablehlo.divide %119, %139 : tensor<256xf32> + %141 = stablehlo.divide %137, %139 : tensor<256xf32> + %142 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x4xf32> + %143 = stablehlo.broadcast_in_dim %140, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %144 = stablehlo.select %117, %143, %101 : tensor<256x4xi1>, tensor<256x4xf32> + %145 = stablehlo.add %142, %144 : tensor<256x4xf32> + %146 = stablehlo.broadcast_in_dim %141, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %147 = stablehlo.select %135, %146, %101 : tensor<256x4xi1>, tensor<256x4xf32> + %148 = stablehlo.add %145, %147 : tensor<256x4xf32> + %149 = stablehlo.broadcast_in_dim %90, dims = [1, 2] : (tensor<256x8xf32>) -> tensor<4x256x8xf32> + %150 = stablehlo.dot_general %149, %arg13, batching_dims = [0] x [0], contracting_dims = [2] x [2] : (tensor<4x256x8xf32>, tensor<4x12x8xf32>) -> tensor<4x256x12xf32> + %151 = stablehlo.dot_general %149, %arg14, batching_dims = [0] x [0], contracting_dims = [2] x [2] : (tensor<4x256x8xf32>, tensor<4x12x8xf32>) -> tensor<4x256x12xf32> + %152 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<4x256x12xf32> + %153 = stablehlo.negate %150 : tensor<4x256x12xf32> + %154 = stablehlo.exponential %153 : tensor<4x256x12xf32> + %155 = stablehlo.add %152, %154 : tensor<4x256x12xf32> + %156 = stablehlo.divide %152, %155 : tensor<4x256x12xf32> + %157 = stablehlo.multiply %150, %156 : tensor<4x256x12xf32> + %158 = stablehlo.multiply %157, %151 : tensor<4x256x12xf32> + %159 = stablehlo.dot_general %158, %arg15, batching_dims = [0] x [0], contracting_dims = [2] x [2] : (tensor<4x256x12xf32>, tensor<4x8x12xf32>) -> tensor<4x256x8xf32> + %160 = stablehlo.dot_general %148, %159, batching_dims = [0] x [1], contracting_dims = [1] x [0] : (tensor<256x4xf32>, tensor<4x256x8xf32>) -> tensor<256x8xf32> + %161 = stablehlo.dot_general %90, %arg16, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<10x8xf32>) -> tensor<256x10xf32> + %162 = stablehlo.dot_general %90, %arg17, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<10x8xf32>) -> tensor<256x10xf32> + %163 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x10xf32> + %164 = stablehlo.negate %161 : tensor<256x10xf32> + %165 = stablehlo.exponential %164 : tensor<256x10xf32> + %166 = stablehlo.add %163, %165 : tensor<256x10xf32> + %167 = stablehlo.divide %163, %166 : tensor<256x10xf32> + %168 = stablehlo.multiply %161, %167 : tensor<256x10xf32> + %169 = stablehlo.multiply %168, %162 : tensor<256x10xf32> + %170 = stablehlo.dot_general %169, %arg18, contracting_dims = [1] x [1] : (tensor<256x10xf32>, tensor<8x10xf32>) -> tensor<256x8xf32> + %171 = stablehlo.dot_general %90, %arg19, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<1x8xf32>) -> tensor<256x1xf32> + %172 = stablehlo.reshape %171 : (tensor<256x1xf32>) -> tensor<256xf32> + %173 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256xf32> + %174 = stablehlo.negate %172 : tensor<256xf32> + %175 = stablehlo.exponential %174 : tensor<256xf32> + %176 = stablehlo.add %173, %175 : tensor<256xf32> + %177 = stablehlo.divide %173, %176 : tensor<256xf32> + %178 = stablehlo.broadcast_in_dim %177, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %179 = stablehlo.multiply %170, %178 : tensor<256x8xf32> + %180 = stablehlo.add %160, %179 : tensor<256x8xf32> + %181 = stablehlo.add %79, %180 : tensor<256x8xf32> + %182 = stablehlo.multiply %181, %181 : tensor<256x8xf32> + %183 = stablehlo.reduce(%182 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %184 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %185 = stablehlo.divide %183, %184 : tensor<256xf32> + %186 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %187 = stablehlo.add %185, %186 : tensor<256xf32> + %188 = stablehlo.rsqrt %187 : tensor<256xf32> + %189 = stablehlo.broadcast_in_dim %188, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %190 = stablehlo.multiply %181, %189 : tensor<256x8xf32> + %191 = stablehlo.broadcast_in_dim %arg20, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %192 = stablehlo.multiply %190, %191 : tensor<256x8xf32> + %193 = stablehlo.dot_general %192, %arg24, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %194 = stablehlo.broadcast_in_dim %arg27, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %195 = stablehlo.add %193, %194 : tensor<256x8xf32> + %196 = stablehlo.reshape %195 : (tensor<256x8xf32>) -> tensor<256x2x4xf32> + %197 = stablehlo.dot_general %192, %arg22, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %198 = stablehlo.broadcast_in_dim %arg26, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %199 = stablehlo.add %197, %198 : tensor<256x4xf32> + %200 = stablehlo.reshape %199 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %201 = stablehlo.dot_general %192, %arg25, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %202 = stablehlo.broadcast_in_dim %arg28, dims = [1] : (tensor<4xf32>) -> tensor<256x4xf32> + %203 = stablehlo.add %201, %202 : tensor<256x4xf32> + %204 = stablehlo.reshape %203 : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %205 = stablehlo.broadcast_in_dim %13, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %206 = stablehlo.broadcast_in_dim %14, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x2x4xf32> + %207 = stablehlo.multiply %196, %205 : tensor<256x2x4xf32> + %208 = stablehlo.slice %196 [0:256, 0:2, 0:2] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %209 = stablehlo.slice %196 [0:256, 0:2, 2:4] : (tensor<256x2x4xf32>) -> tensor<256x2x2xf32> + %210 = stablehlo.negate %209 : tensor<256x2x2xf32> + %211 = stablehlo.concatenate %210, %208, dim = 2 : (tensor<256x2x2xf32>, tensor<256x2x2xf32>) -> tensor<256x2x4xf32> + %212 = stablehlo.multiply %211, %206 : tensor<256x2x4xf32> + %213 = stablehlo.add %207, %212 : tensor<256x2x4xf32> + %214 = stablehlo.broadcast_in_dim %13, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %215 = stablehlo.broadcast_in_dim %14, dims = [0, 2] : (tensor<256x4xf32>) -> tensor<256x1x4xf32> + %216 = stablehlo.multiply %200, %214 : tensor<256x1x4xf32> + %217 = stablehlo.slice %200 [0:256, 0:1, 0:2] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %218 = stablehlo.slice %200 [0:256, 0:1, 2:4] : (tensor<256x1x4xf32>) -> tensor<256x1x2xf32> + %219 = stablehlo.negate %218 : tensor<256x1x2xf32> + %220 = stablehlo.concatenate %219, %217, dim = 2 : (tensor<256x1x2xf32>, tensor<256x1x2xf32>) -> tensor<256x1x4xf32> + %221 = stablehlo.multiply %220, %215 : tensor<256x1x4xf32> + %222 = stablehlo.add %216, %221 : tensor<256x1x4xf32> + %223 = stablehlo.reshape %222 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %224 = stablehlo.dynamic_update_slice %59, %223, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %225 = stablehlo.reshape %204 : (tensor<256x1x4xf32>) -> tensor<1x256x1x4xf32> + %226 = stablehlo.dynamic_update_slice %61, %225, %11, %9, %9, %9 : (tensor<2x256x1x4xf32>, tensor<1x256x1x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x256x1x4xf32> + %227 = stablehlo.reshape %213 : (tensor<256x2x4xf32>) -> tensor<256x1x2x4xf32> + %228 = stablehlo.dot_general %227, %222, batching_dims = [1] x [1], contracting_dims = [3] x [2] : (tensor<256x1x2x4xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x256xf32> + %229 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor) -> tensor<1x256x2x256xf32> + %230 = stablehlo.multiply %228, %229 : tensor<1x256x2x256xf32> + %231 = stablehlo.broadcast_in_dim %arg40, dims = [1, 3] : (tensor<256x256xf32>) -> tensor<1x256x2x256xf32> + %232 = stablehlo.add %230, %231 : tensor<1x256x2x256xf32> + %233 = stablehlo.reduce(%232 init: %4) applies stablehlo.maximum across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %234 = stablehlo.broadcast_in_dim %233, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %235 = stablehlo.subtract %232, %234 : tensor<1x256x2x256xf32> + %236 = stablehlo.exponential %235 : tensor<1x256x2x256xf32> + %237 = stablehlo.reduce(%236 init: %2) applies stablehlo.add across dimensions = [3] : (tensor<1x256x2x256xf32>, tensor) -> tensor<1x256x2xf32> + %238 = stablehlo.broadcast_in_dim %237, dims = [0, 1, 2] : (tensor<1x256x2xf32>) -> tensor<1x256x2x256xf32> + %239 = stablehlo.divide %236, %238 : tensor<1x256x2x256xf32> + %240 = stablehlo.dot_general %239, %204, batching_dims = [0] x [1], contracting_dims = [3] x [0] : (tensor<1x256x2x256xf32>, tensor<256x1x4xf32>) -> tensor<1x256x2x4xf32> + %241 = stablehlo.transpose %240, dims = [1, 0, 2, 3] : (tensor<1x256x2x4xf32>) -> tensor<256x1x2x4xf32> + %242 = stablehlo.reshape %241 : (tensor<256x1x2x4xf32>) -> tensor<256x8xf32> + %243 = stablehlo.dot_general %242, %arg23, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<8x8xf32>) -> tensor<256x8xf32> + %244 = stablehlo.add %181, %243 : tensor<256x8xf32> + %245 = stablehlo.multiply %244, %244 : tensor<256x8xf32> + %246 = stablehlo.reduce(%245 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %247 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %248 = stablehlo.divide %246, %247 : tensor<256xf32> + %249 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %250 = stablehlo.add %248, %249 : tensor<256xf32> + %251 = stablehlo.rsqrt %250 : tensor<256xf32> + %252 = stablehlo.broadcast_in_dim %251, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %253 = stablehlo.multiply %244, %252 : tensor<256x8xf32> + %254 = stablehlo.broadcast_in_dim %arg21, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %255 = stablehlo.multiply %253, %254 : tensor<256x8xf32> + %256 = stablehlo.dot_general %255, %arg29, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<4x8xf32>) -> tensor<256x4xf32> + %257 = stablehlo.reduce(%256 init: %4) applies stablehlo.maximum across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %258 = stablehlo.broadcast_in_dim %257, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %259 = stablehlo.subtract %256, %258 : tensor<256x4xf32> + %260 = stablehlo.exponential %259 : tensor<256x4xf32> + %261 = stablehlo.reduce(%260 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %262 = stablehlo.broadcast_in_dim %261, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %263 = stablehlo.divide %260, %262 : tensor<256x4xf32> + %264 = stablehlo.iota dim = 0 : tensor<4xi32> + %265 = stablehlo.broadcast_in_dim %264, dims = [1] : (tensor<4xi32>) -> tensor<256x4xi32> + %266 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x4xf32> + %267 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor) -> tensor<256x4xf32> + %268 = stablehlo.constant dense<0> : tensor + %269 = stablehlo.constant dense<0xFF800000> : tensor + %270 = stablehlo.iota dim = 1 : tensor<256x4xi32> + %271:2 = stablehlo.reduce(%263 init: %269), (%270 init: %268) across dimensions = [1] : (tensor<256x4xf32>, tensor<256x4xi32>, tensor, tensor) -> (tensor<256xf32>, tensor<256xi32>) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %272 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %273 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %274 = stablehlo.or %272, %273 : tensor + %275 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %276 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %277 = stablehlo.and %275, %276 : tensor + %278 = stablehlo.or %274, %277 : tensor + %279 = stablehlo.select %274, %amv_l, %amv_r : tensor, tensor + %280 = stablehlo.select %278, %ami_l, %ami_r : tensor, tensor + stablehlo.return %279, %280 : tensor, tensor + } + %281 = stablehlo.broadcast_in_dim %271#1, dims = [0] : (tensor<256xi32>) -> tensor<256x4xi32> + %282 = stablehlo.compare EQ, %265, %281, SIGNED : (tensor<256x4xi32>, tensor<256x4xi32>) -> tensor<256x4xi1> + %283 = stablehlo.select %282, %263, %266 : tensor<256x4xi1>, tensor<256x4xf32> + %284 = stablehlo.reduce(%283 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %285 = stablehlo.select %282, %267, %263 : tensor<256x4xi1>, tensor<256x4xf32> + %286 = stablehlo.constant dense<0> : tensor + %287 = stablehlo.constant dense<0xFF800000> : tensor + %288 = stablehlo.iota dim = 1 : tensor<256x4xi32> + %289:2 = stablehlo.reduce(%285 init: %287), (%288 init: %286) across dimensions = [1] : (tensor<256x4xf32>, tensor<256x4xi32>, tensor, tensor) -> (tensor<256xf32>, tensor<256xi32>) + reducer(%amv_l: tensor, %amv_r: tensor) (%ami_l: tensor, %ami_r: tensor) { + %290 = stablehlo.compare GT, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %291 = stablehlo.compare NE, %amv_l, %amv_l, FLOAT : (tensor, tensor) -> tensor + %292 = stablehlo.or %290, %291 : tensor + %293 = stablehlo.compare EQ, %amv_l, %amv_r, FLOAT : (tensor, tensor) -> tensor + %294 = stablehlo.compare LT, %ami_l, %ami_r, SIGNED : (tensor, tensor) -> tensor + %295 = stablehlo.and %293, %294 : tensor + %296 = stablehlo.or %292, %295 : tensor + %297 = stablehlo.select %292, %amv_l, %amv_r : tensor, tensor + %298 = stablehlo.select %296, %ami_l, %ami_r : tensor, tensor + stablehlo.return %297, %298 : tensor, tensor + } + %299 = stablehlo.broadcast_in_dim %289#1, dims = [0] : (tensor<256xi32>) -> tensor<256x4xi32> + %300 = stablehlo.compare EQ, %265, %299, SIGNED : (tensor<256x4xi32>, tensor<256x4xi32>) -> tensor<256x4xi1> + %301 = stablehlo.select %300, %285, %266 : tensor<256x4xi1>, tensor<256x4xf32> + %302 = stablehlo.reduce(%301 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x4xf32>, tensor) -> tensor<256xf32> + %303 = stablehlo.select %300, %267, %285 : tensor<256x4xi1>, tensor<256x4xf32> + %304 = stablehlo.add %284, %302 : tensor<256xf32> + %305 = stablehlo.divide %284, %304 : tensor<256xf32> + %306 = stablehlo.divide %302, %304 : tensor<256xf32> + %307 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor) -> tensor<256x4xf32> + %308 = stablehlo.broadcast_in_dim %305, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %309 = stablehlo.select %282, %308, %266 : tensor<256x4xi1>, tensor<256x4xf32> + %310 = stablehlo.add %307, %309 : tensor<256x4xf32> + %311 = stablehlo.broadcast_in_dim %306, dims = [0] : (tensor<256xf32>) -> tensor<256x4xf32> + %312 = stablehlo.select %300, %311, %266 : tensor<256x4xi1>, tensor<256x4xf32> + %313 = stablehlo.add %310, %312 : tensor<256x4xf32> + %314 = stablehlo.broadcast_in_dim %255, dims = [1, 2] : (tensor<256x8xf32>) -> tensor<4x256x8xf32> + %315 = stablehlo.dot_general %314, %arg30, batching_dims = [0] x [0], contracting_dims = [2] x [2] : (tensor<4x256x8xf32>, tensor<4x12x8xf32>) -> tensor<4x256x12xf32> + %316 = stablehlo.dot_general %314, %arg31, batching_dims = [0] x [0], contracting_dims = [2] x [2] : (tensor<4x256x8xf32>, tensor<4x12x8xf32>) -> tensor<4x256x12xf32> + %317 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<4x256x12xf32> + %318 = stablehlo.negate %315 : tensor<4x256x12xf32> + %319 = stablehlo.exponential %318 : tensor<4x256x12xf32> + %320 = stablehlo.add %317, %319 : tensor<4x256x12xf32> + %321 = stablehlo.divide %317, %320 : tensor<4x256x12xf32> + %322 = stablehlo.multiply %315, %321 : tensor<4x256x12xf32> + %323 = stablehlo.multiply %322, %316 : tensor<4x256x12xf32> + %324 = stablehlo.dot_general %323, %arg32, batching_dims = [0] x [0], contracting_dims = [2] x [2] : (tensor<4x256x12xf32>, tensor<4x8x12xf32>) -> tensor<4x256x8xf32> + %325 = stablehlo.dot_general %313, %324, batching_dims = [0] x [1], contracting_dims = [1] x [0] : (tensor<256x4xf32>, tensor<4x256x8xf32>) -> tensor<256x8xf32> + %326 = stablehlo.dot_general %255, %arg33, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<10x8xf32>) -> tensor<256x10xf32> + %327 = stablehlo.dot_general %255, %arg34, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<10x8xf32>) -> tensor<256x10xf32> + %328 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256x10xf32> + %329 = stablehlo.negate %326 : tensor<256x10xf32> + %330 = stablehlo.exponential %329 : tensor<256x10xf32> + %331 = stablehlo.add %328, %330 : tensor<256x10xf32> + %332 = stablehlo.divide %328, %331 : tensor<256x10xf32> + %333 = stablehlo.multiply %326, %332 : tensor<256x10xf32> + %334 = stablehlo.multiply %333, %327 : tensor<256x10xf32> + %335 = stablehlo.dot_general %334, %arg35, contracting_dims = [1] x [1] : (tensor<256x10xf32>, tensor<8x10xf32>) -> tensor<256x8xf32> + %336 = stablehlo.dot_general %255, %arg36, contracting_dims = [1] x [1] : (tensor<256x8xf32>, tensor<1x8xf32>) -> tensor<256x1xf32> + %337 = stablehlo.reshape %336 : (tensor<256x1xf32>) -> tensor<256xf32> + %338 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor) -> tensor<256xf32> + %339 = stablehlo.negate %337 : tensor<256xf32> + %340 = stablehlo.exponential %339 : tensor<256xf32> + %341 = stablehlo.add %338, %340 : tensor<256xf32> + %342 = stablehlo.divide %338, %341 : tensor<256xf32> + %343 = stablehlo.broadcast_in_dim %342, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %344 = stablehlo.multiply %335, %343 : tensor<256x8xf32> + %345 = stablehlo.add %325, %344 : tensor<256x8xf32> + %346 = stablehlo.add %244, %345 : tensor<256x8xf32> + %347 = stablehlo.multiply %346, %346 : tensor<256x8xf32> + %348 = stablehlo.reduce(%347 init: %2) applies stablehlo.add across dimensions = [1] : (tensor<256x8xf32>, tensor) -> tensor<256xf32> + %349 = stablehlo.broadcast_in_dim %7, dims = [] : (tensor) -> tensor<256xf32> + %350 = stablehlo.divide %348, %349 : tensor<256xf32> + %351 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor) -> tensor<256xf32> + %352 = stablehlo.add %350, %351 : tensor<256xf32> + %353 = stablehlo.rsqrt %352 : tensor<256xf32> + %354 = stablehlo.broadcast_in_dim %353, dims = [0] : (tensor<256xf32>) -> tensor<256x8xf32> + %355 = stablehlo.multiply %346, %354 : tensor<256x8xf32> + %356 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<8xf32>) -> tensor<256x8xf32> + %357 = stablehlo.multiply %355, %356 : tensor<256x8xf32> + %358 = stablehlo.constant dense<1> : tensor + %359 = stablehlo.subtract %arg39, %358 : tensor + %360 = stablehlo.dynamic_slice %357, %359, %9, sizes = [1, 8] : (tensor<256x8xf32>, tensor, tensor) -> tensor<1x8xf32> + %361 = stablehlo.reshape %360 : (tensor<1x8xf32>) -> tensor<8xf32> + %362 = stablehlo.dot_general %361, %arg2, contracting_dims = [0] x [1] : (tensor<8xf32>, tensor<10x8xf32>) -> tensor<10xf32> + return %362, %224, %226 : tensor<10xf32>, tensor<2x256x1x4xf32>, tensor<2x256x1x4xf32> + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/context_tests.rs b/src/lib/mlxcel-xla/src/emitter/context_tests.rs index 7206f9e1..76b7cef6 100644 --- a/src/lib/mlxcel-xla/src/emitter/context_tests.rs +++ b/src/lib/mlxcel-xla/src/emitter/context_tests.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::{Config, emit_decode, emit_decode_ragged, emit_prefill}; +use super::{Config, emit_decode, emit_decode_ragged, emit_prefill, emit_prefill_embeddings}; fn tiny_config(context_capacity: usize) -> Config { Config { @@ -31,11 +31,13 @@ fn tiny_config(context_capacity: usize) -> Config { fn assert_capacity_shapes(context_capacity: usize) { let cfg = tiny_config(context_capacity); let prefill = emit_prefill(&cfg, false); + let embeddings_prefill = emit_prefill_embeddings(&cfg, false); let decode = emit_decode(&cfg, false); let ragged = emit_decode_ragged(&cfg, 4, false); let single_cache = format!("tensor<2x{context_capacity}x1x4xf32>"); let batch_cache = format!("tensor<4x2x{context_capacity}x1x4xf32>"); let prompt = format!("tensor<{context_capacity}xi32>"); + let embeddings = format!("tensor<{context_capacity}x8xf32>"); let causal_mask = format!("tensor<{context_capacity}x{context_capacity}xf32>"); let ragged_mask = format!("tensor<4x{context_capacity}xf32>"); let rope = format!("tensor<{context_capacity}x4xf32>"); @@ -56,6 +58,26 @@ fn assert_capacity_shapes(context_capacity: usize) { prefill.contains(&rope), "prefill position table uses {rope}" ); + assert!( + embeddings_prefill.contains(&embeddings), + "embeddings prefill input uses {embeddings}" + ); + assert!( + embeddings_prefill.matches(&prompt).count() >= 1, + "embeddings prefill positions use {prompt}" + ); + assert!( + embeddings_prefill.matches(&single_cache).count() >= 3, + "embeddings prefill returns matching K/V" + ); + assert!( + embeddings_prefill.contains(&causal_mask), + "embeddings prefill bias uses {causal_mask}" + ); + assert!( + embeddings_prefill.contains(&rope), + "embeddings prefill position table uses {rope}" + ); assert!( decode.matches(&single_cache).count() >= 5, diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index 3283913c..2c5317fd 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -70,8 +70,11 @@ pub(crate) use builder::{ }; #[allow(unused_imports)] pub(crate) use model::{ + PREFILL_EMBEDDINGS_MASKED_VALUE, PrefillEmbeddingsDType, PrefillEmbeddingsInputMetadata, emit_decode, emit_decode_batched, emit_decode_ragged, emit_decode_ragged_with, - emit_decode_with, emit_moe_probe, emit_moe_probe_with, emit_prefill, emit_prefill_with, + emit_decode_with, emit_moe_probe, emit_moe_probe_with, emit_prefill, emit_prefill_embeddings, + emit_prefill_embeddings_with, emit_prefill_with, validate_prefill_embeddings_attention_bias, + validate_prefill_embeddings_metadata, }; #[cfg(test)] @@ -736,6 +739,170 @@ mod tests { .expect("write decode_logits.mlir"); } + /// Opt-in graph dump for the real IREE token-vs-embeddings parity fixture. + /// Both modules are emitted from the exact same config and precision so the + /// fixture can feed one weight set to each and compare logits plus every K/V + /// element. This stays pure Rust; the Python harness owns IREE compilation and + /// execution. + #[test] + #[ignore = "opt-in: writes token and embeddings prefill graphs for IREE parity"] + fn dump_prefill_embeddings_parity_graphs() { + let cfg_path = std::env::var("MLXCEL_DUMP_CONFIG") + .expect("set MLXCEL_DUMP_CONFIG to a config.json path"); + let dir = std::env::var("MLXCEL_DUMP_DIR").expect("set MLXCEL_DUMP_DIR"); + let text = std::fs::read_to_string(&cfg_path).expect("read MLXCEL_DUMP_CONFIG"); + let cfg = Config::from_json_str(&text).expect("parse config.json"); + let dir = std::path::Path::new(&dir); + std::fs::write(dir.join("prefill_logits.mlir"), emit_prefill(&cfg, false)) + .expect("write token prefill graph"); + std::fs::write( + dir.join("prefill_embeddings_logits.mlir"), + emit_prefill_embeddings(&cfg, false), + ) + .expect("write embeddings prefill graph"); + } + + #[test] + fn prefill_embeddings_schema_bypasses_lookup_and_embedding_scale() { + let cfg = gemma2_like(None); + let lp = cfg.context_capacity; + let token = emit_prefill(&cfg, false); + let embeddings = emit_prefill_embeddings(&cfg, false); + let sampled_embeddings = emit_prefill_embeddings(&cfg, true); + + assert!(token.starts_with("module @prefill {")); + assert!(embeddings.starts_with("module @prefill_embeddings {")); + assert!(embeddings.contains(&format!("tensor<{lp}x8xf32> loc(\"embeddings\")"))); + assert!(embeddings.contains(&format!("tensor<{lp}xi32> loc(\"positions\")"))); + assert!(embeddings.contains("tensor loc(\"real_len\")")); + assert!(embeddings.contains(&format!("tensor<{lp}x{lp}xf32> loc(\"attention_bias\")"))); + assert!(!embeddings.contains("loc(\"tokens\")")); + assert!( + sampled_embeddings.contains(&format!( + ") -> (tensor, tensor<4x{lp}x1x6xf32>, tensor<4x{lp}x1x6xf32>)" + )), + "sampled embeddings prefill keeps the scalar argmax plus K/V return contract" + ); + + let ordered = [ + "loc(\"embeddings\")", + "loc(\"positions\")", + "loc(\"real_len\")", + "loc(\"attention_bias\")", + ]; + let offsets: Vec<_> = ordered + .iter() + .map(|needle| embeddings.find(needle).expect("schema arg present")) + .collect(); + assert!( + offsets.windows(2).all(|w| w[0] < w[1]), + "runtime inputs keep their documented deterministic order" + ); + + assert!( + token.contains("\"stablehlo.gather\"(%arg0,"), + "token prefill gathers the embedding table" + ); + assert!( + !embeddings.contains("\"stablehlo.gather\"(%arg0,"), + "embeddings prefill must not gather the embedding table" + ); + let scale_hex = format!("0x{:08X}", cfg.embed_normalizer().to_bits()); + assert!( + token.contains(&scale_hex), + "token prefill emits Gemma scaling" + ); + assert!( + !embeddings.contains(&scale_hex), + "post-scale embeddings must not be scaled a second time" + ); + } + + #[test] + fn prefill_embeddings_preserves_tied_and_untied_head_weights() { + let tied = qwen_like(false); + let tied_mlir = emit_prefill_embeddings(&tied, false); + assert!(tied_mlir.contains("loc(\"params['embed']\")")); + assert!(!tied_mlir.contains("params['lm_head']")); + let tied_tail = tied_mlir + .lines() + .rfind(|line| line.contains("stablehlo.dot_general")) + .expect("tied logits dot"); + assert!(tied_tail.contains("%arg0"), "tied logits reuse embed"); + + let mut untied = tied; + untied.tie_word_embeddings = false; + let untied_mlir = emit_prefill_embeddings(&untied, false); + assert!(untied_mlir.contains("loc(\"params['embed']\")")); + assert!(untied_mlir.contains("loc(\"params['lm_head']\")")); + let untied_tail = untied_mlir + .lines() + .rfind(|line| line.contains("stablehlo.dot_general")) + .expect("untied logits dot"); + assert!(untied_tail.contains("%arg2"), "untied logits use lm_head"); + } + + #[test] + fn prefill_embeddings_metadata_and_bias_fail_closed() { + let cfg = qwen_like(false); + let lp = cfg.context_capacity; + let valid = PrefillEmbeddingsInputMetadata::canonical(&cfg); + validate_prefill_embeddings_metadata(&cfg, &valid).expect("canonical metadata"); + + let invalid = [ + PrefillEmbeddingsInputMetadata { + embeddings_shape: [lp - 1, cfg.hidden], + ..valid + }, + PrefillEmbeddingsInputMetadata { + embeddings_shape: [lp, cfg.hidden + 1], + ..valid + }, + PrefillEmbeddingsInputMetadata { + embeddings_dtype: PrefillEmbeddingsDType::F16, + ..valid + }, + PrefillEmbeddingsInputMetadata { + positions_shape: [lp - 1], + ..valid + }, + PrefillEmbeddingsInputMetadata { + attention_bias_shape: [lp, lp - 1], + ..valid + }, + PrefillEmbeddingsInputMetadata { + attention_bias_dtype: PrefillEmbeddingsDType::Bf16, + ..valid + }, + PrefillEmbeddingsInputMetadata { + real_len: 0, + ..valid + }, + PrefillEmbeddingsInputMetadata { + real_len: lp + 1, + ..valid + }, + ]; + for metadata in invalid { + assert!( + validate_prefill_embeddings_metadata(&cfg, &metadata).is_err(), + "invalid metadata must fail: {metadata:?}" + ); + } + + let mut bias = vec![PREFILL_EMBEDDINGS_MASKED_VALUE; lp * lp]; + bias[0] = 0.0; + validate_prefill_embeddings_attention_bias(&cfg, &bias).expect("canonical finite bias"); + for invalid_value in [f32::NEG_INFINITY, f32::NAN, -1.0] { + bias[17] = invalid_value; + assert!( + validate_prefill_embeddings_attention_bias(&cfg, &bias).is_err(), + "invalid bias value {invalid_value} must fail" + ); + } + assert!(validate_prefill_embeddings_attention_bias(&cfg, &bias[..bias.len() - 1]).is_err()); + } + /// Plain RoPE base frequencies are the textbook `1 / theta^(2i/head_dim)` /// (Qwen2), distinct from the llama3-scaled table. #[test] diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index 07fd9bd1..1e3df2ad 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/src/lib/mlxcel-xla/src/emitter/model.rs @@ -24,6 +24,144 @@ use super::config::{Config, NormStyle}; use super::moe::{self, MoeLayerW, MoeSharedW}; use super::rope; +/// Canonical finite additive-attention value for a disallowed prefill edge. +/// +/// The embeddings entry accepts only `0.0` (allowed) and this value (masked). +/// Negative infinity is deliberately rejected so every caller, compiler target, +/// and softmax lowering observes the same finite input convention as the existing +/// token prefill graph. +pub(crate) const PREFILL_EMBEDDINGS_MASKED_VALUE: f32 = -1e30; + +/// Runtime tensor element types understood by the prefill-embeddings schema +/// validator. The graph currently consumes f32 hidden states and an f32 additive +/// bias; the other variants exist so callers get an actionable error before IREE +/// compilation/invocation instead of silently reinterpreting bytes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum PrefillEmbeddingsDType { + F32, + F16, + Bf16, +} + +/// Shape/dtype metadata validated before compiling or invoking +/// `prefill_embeddings.main`. +/// +/// Argument order after the model weights is stable and deliberately documented +/// here: `embeddings`, `positions`, `real_len`, `attention_bias`. Positions remain +/// one-dimensional in this entry; a later M-RoPE entry can add a distinct position +/// representation without overloading this contract. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct PrefillEmbeddingsInputMetadata { + pub embeddings_shape: [usize; 2], + pub embeddings_dtype: PrefillEmbeddingsDType, + pub positions_shape: [usize; 1], + pub attention_bias_shape: [usize; 2], + pub attention_bias_dtype: PrefillEmbeddingsDType, + pub real_len: usize, +} + +impl PrefillEmbeddingsInputMetadata { + /// Metadata for the static context-capacity schema emitted for `c`. + pub(crate) fn canonical(c: &Config) -> Self { + let lp = c.context_capacity; + Self { + embeddings_shape: [lp, c.hidden], + embeddings_dtype: PrefillEmbeddingsDType::F32, + positions_shape: [lp], + attention_bias_shape: [lp, lp], + attention_bias_dtype: PrefillEmbeddingsDType::F32, + real_len: lp, + } + } +} + +/// Validate the static prefill-from-embeddings input contract before compiling or +/// invoking the graph. +/// +/// `real_len` selects the output row but does not alter padded-row computation: +/// callers provide a complete `[Lp, Lp]` bias for the whole static bucket. For +/// token-parity padding, padded rows therefore retain the same causal pattern as +/// ordinary rows, while real rows mask every future/padded key. +pub(crate) fn validate_prefill_embeddings_metadata( + c: &Config, + metadata: &PrefillEmbeddingsInputMetadata, +) -> Result<(), String> { + let lp = c.context_capacity; + let want_embeddings = [lp, c.hidden]; + if metadata.embeddings_shape != want_embeddings { + return Err(format!( + "prefill embeddings shape {:?} does not match required {:?}", + metadata.embeddings_shape, want_embeddings + )); + } + if metadata.embeddings_dtype != PrefillEmbeddingsDType::F32 { + return Err(format!( + "prefill embeddings dtype {:?} is unsupported; expected F32", + metadata.embeddings_dtype + )); + } + if metadata.positions_shape != [lp] { + return Err(format!( + "prefill positions shape {:?} does not match required [{}]", + metadata.positions_shape, lp + )); + } + if metadata.attention_bias_shape != [lp, lp] { + return Err(format!( + "prefill attention-bias shape {:?} does not match required [{}, {}]", + metadata.attention_bias_shape, lp, lp + )); + } + if metadata.attention_bias_dtype != PrefillEmbeddingsDType::F32 { + return Err(format!( + "prefill attention-bias dtype {:?} is unsupported; expected F32", + metadata.attention_bias_dtype + )); + } + if !(1..=lp).contains(&metadata.real_len) { + return Err(format!( + "prefill real_len {} is outside the required range 1..={lp}", + metadata.real_len, + )); + } + Ok(()) +} + +/// Validate a canonical additive attention-bias payload. Only `0.0` (allowed) +/// and [`PREFILL_EMBEDDINGS_MASKED_VALUE`] (masked) are accepted; NaN, infinity, +/// and intermediate additive values are rejected to keep polarity unambiguous. +pub(crate) fn validate_prefill_embeddings_attention_bias( + c: &Config, + bias: &[f32], +) -> Result<(), String> { + let expected = c + .context_capacity + .checked_mul(c.context_capacity) + .ok_or_else(|| { + format!( + "prefill attention-bias size overflows for context capacity {}", + c.context_capacity + ) + })?; + if bias.len() != expected { + return Err(format!( + "prefill attention bias has {} elements; expected {expected}", + bias.len() + )); + } + if let Some((index, value)) = bias + .iter() + .copied() + .enumerate() + .find(|(_, value)| *value != 0.0 && *value != PREFILL_EMBEDDINGS_MASKED_VALUE) + { + return Err(format!( + "prefill attention bias at flat index {index} is {value}; expected only 0 or {}", + PREFILL_EMBEDDINGS_MASKED_VALUE + )); + } + Ok(()) +} /// Per-layer weight handles (JAX alphabetical order: down, gate, in_ln, /// post_ln, up, wk, wo, wq, wv). `bk`/`bq`/`bv` are the q/k/v projection biases, /// present only for an architecture with `qkv_bias` (Qwen2); `None` for Llama, @@ -2341,20 +2479,43 @@ pub fn emit_decode_ragged_with( // processed at once over an [Lp] sequence axis with an [Lp,Lp] causal mask, and // the returned logit is the row at real_len-1 (the last real prompt token). -/// Prefill arg handles. Weights are identical to decode (same order/locs); the -/// trailing inputs are tokens/positions/real_len with no caches. +/// Which stable prefill entry schema to build. The token entry remains unchanged; +/// the embeddings entry replaces `tokens` with post-scale hidden states and adds +/// an explicit additive attention bias. +#[derive(Clone, Copy)] +enum PrefillInputKind { + Tokens, + Embeddings, +} + +enum PrefillInput { + Tokens(Val), + Embeddings { hidden: Val, attention_bias: Val }, +} + +/// Prefill arg handles. Weights are identical to decode (same order/locs), which +/// intentionally retains `params['embed']`: token prefill gathers from it and a +/// tied embeddings prefill reuses it as the LM head. The trailing schema is: +/// +/// - token: `tokens`, `positions`, `real_len`; +/// - embeddings: `embeddings`, `positions`, `real_len`, `attention_bias`. struct PrefillArgs { embed: Val, final_norm: Val, final_norm_bias: Option, lm_head: Option, layers: Vec, - tokens: Val, + input: PrefillInput, positions: Val, real_len: Val, } -fn build_prefill_arg_schema(b: &mut Builder, c: &Config, lp: usize) -> (Vec, PrefillArgs) { +fn build_prefill_arg_schema( + b: &mut Builder, + c: &Config, + lp: usize, + input_kind: PrefillInputKind, +) -> (Vec, PrefillArgs) { let h = c.hidden; let v = c.vocab; @@ -2381,12 +2542,26 @@ fn build_prefill_arg_schema(b: &mut Builder, c: &Config, lp: usize) -> (Vec ( + Some(take_arg( + &mut decls, + &mut idx, + Ty::new(vec![lp], "i32"), + "tokens".into(), + )), + None, + ), + PrefillInputKind::Embeddings => ( + None, + Some(take_arg( + &mut decls, + &mut idx, + Ty::f32(vec![lp, h]), + "embeddings".into(), + )), + ), + }; let positions = take_arg( &mut decls, &mut idx, @@ -2394,6 +2569,19 @@ fn build_prefill_arg_schema(b: &mut Builder, c: &Config, lp: usize) -> (Vec PrefillInput::Embeddings { + hidden, + attention_bias: take_arg( + &mut decls, + &mut idx, + Ty::f32(vec![lp, lp]), + "attention_bias".into(), + ), + }, + (Some(tokens), None) => PrefillInput::Tokens(tokens), + _ => unreachable!("each prefill schema has exactly one input representation"), + }; ( decls, @@ -2403,7 +2591,7 @@ fn build_prefill_arg_schema(b: &mut Builder, c: &Config, lp: usize) -> (Vec String { } pub fn emit_prefill_with(c: &Config, sample: bool, precision: Precision) -> String { + emit_prefill_module(c, sample, precision, PrefillInputKind::Tokens) +} + +/// Emit the distinct prefill-from-embeddings StableHLO module at the ambient +/// precision. The input hidden states are already post-token-embedding-scale; +/// this entry performs neither a token gather nor [`scale_embedding`]. +pub(crate) fn emit_prefill_embeddings(c: &Config, sample: bool) -> String { + emit_prefill_embeddings_with(c, sample, precision_from_env()) +} + +/// Emit `prefill_embeddings.main` at an explicit contraction precision. +pub(crate) fn emit_prefill_embeddings_with( + c: &Config, + sample: bool, + precision: Precision, +) -> String { + emit_prefill_module(c, sample, precision, PrefillInputKind::Embeddings) +} + +fn emit_prefill_module( + c: &Config, + sample: bool, + precision: Precision, + input_kind: PrefillInputKind, +) -> String { let lp = c.context_capacity; let mut b = Builder::new().with_precision(precision); - let (decls, a) = build_prefill_arg_schema(&mut b, c, lp); + let (decls, a) = build_prefill_arg_schema(&mut b, c, lp, input_kind); let k = emit_consts(&mut b, c); let h = c.hidden; let d = c.head_dim; let nkv = c.n_kv; - // --- head: embed gather, per-position rope vectors, [Lp,Lp] causal mask --- - let tok_idx = b.reshape(&a.tokens, vec![lp, 1]); - let emb = b.gather(&a.embed, &tok_idx); // [Lp, H] - let mut x = scale_embedding(&mut b, c, emb, vec![lp, h]); + // --- input head --- + // Token prefill gathers + applies the architecture's embedding scale. The + // embeddings entry consumes already-merged, post-scale hidden states as-is. + let mut x = match &a.input { + PrefillInput::Tokens(tokens) => { + let tok_idx = b.reshape(tokens, vec![lp, 1]); + let emb = b.gather(&a.embed, &tok_idx); // [Lp, H] + scale_embedding(&mut b, c, emb, vec![lp, h]) + } + PrefillInput::Embeddings { hidden, .. } => hidden.clone(), + }; + // --- shared position preparation --- let pos_idx = b.reshape(&a.positions, vec![lp, 1]); let cos = b.gather(&k.cos_table, &pos_idx); // [Lp, d] let sin = b.gather(&k.sin_table, &pos_idx); // [Lp, d] @@ -2516,28 +2737,55 @@ pub fn emit_prefill_with(c: &Config, sample: bool, precision: Precision) -> Stri _ => (None, None), }; - // causal mask [Lp, Lp]: query i attends key j iff j <= i -> additive 0/-1e30 - let irow = b.iota(lp); - let row = b.broadcast(&irow, &[0], vec![lp, lp]); // entry[i,j] = i - let jcol = b.iota(lp); - let col = b.broadcast(&jcol, &[1], vec![lp, lp]); // entry[i,j] = j - let le = b.compare("LE", &col, &row, "SIGNED"); // j <= i - let zeros = b.broadcast(&k.zero, &[], vec![lp, lp]); - let negs = b.broadcast(&k.neg_big, &[], vec![lp, lp]); - let cmask = b.select(&le, &zeros, &negs); // [Lp, Lp] - // Gemma2 local (sliding-window) mask (issue #495): a local layer keeps key j - // for query i iff `i - j < W` (prefill positions are 0..Lp, so the buffer - // index equals the position). And it into the causal `cmask` with one more - // `select`. Emitted only for a windowed config (Gemma2); reused by every local - // layer. No-op on the value when `W >= Lp` (`i - j <= Lp-1 < W`), so a - // short-prompt Gemma2 prefill is unchanged. - let cmask_local = c.sliding_window.map(|w| { - let wc = b.const_i32(w as i32); - let wb = b.broadcast(&wc, &[], vec![lp, lp]); - let age = b.subtract(&row, &col); // i - j - let within = b.compare("LT", &age, &wb, "SIGNED"); - b.select(&within, &cmask, &negs) - }); + // --- attention-bias preparation --- + // Token prefill retains the exact internal causal-mask construction. The + // embeddings entry uses the caller's canonical f32 additive bias directly. + // For a sliding architecture, the local-layer window is intersected with + // either base mask, preserving the architecture invariant without changing + // multimodal/global visibility. + let (cmask, cmask_local) = match &a.input { + PrefillInput::Tokens(_) => { + // causal mask [Lp, Lp]: query i attends key j iff j <= i -> additive 0/-1e30 + let irow = b.iota(lp); + let row = b.broadcast(&irow, &[0], vec![lp, lp]); // entry[i,j] = i + let jcol = b.iota(lp); + let col = b.broadcast(&jcol, &[1], vec![lp, lp]); // entry[i,j] = j + let le = b.compare("LE", &col, &row, "SIGNED"); // j <= i + let zeros = b.broadcast(&k.zero, &[], vec![lp, lp]); + let negs = b.broadcast(&k.neg_big, &[], vec![lp, lp]); + let cmask = b.select(&le, &zeros, &negs); // [Lp, Lp] + // Gemma2 local (sliding-window) mask (issue #495): a local layer keeps key j + // for query i iff `i - j < W` (prefill positions are 0..Lp, so the buffer + // index equals the position). And it into the causal `cmask` with one more + // `select`. Emitted only for a windowed config (Gemma2); reused by every local + // layer. No-op on the value when `W >= Lp` (`i - j <= Lp-1 < W`), so a + // short-prompt Gemma2 prefill is unchanged. + let cmask_local = c.sliding_window.map(|w| { + let wc = b.const_i32(w as i32); + let wb = b.broadcast(&wc, &[], vec![lp, lp]); + let age = b.subtract(&row, &col); // i - j + let within = b.compare("LT", &age, &wb, "SIGNED"); + b.select(&within, &cmask, &negs) + }); + (cmask, cmask_local) + } + PrefillInput::Embeddings { attention_bias, .. } => { + let cmask = attention_bias.clone(); + let cmask_local = c.sliding_window.map(|w| { + let irow = b.iota(lp); + let row = b.broadcast(&irow, &[0], vec![lp, lp]); + let jcol = b.iota(lp); + let col = b.broadcast(&jcol, &[1], vec![lp, lp]); + let wc = b.const_i32(w as i32); + let wb = b.broadcast(&wc, &[], vec![lp, lp]); + let age = b.subtract(&row, &col); + let within = b.compare("LT", &age, &wb, "SIGNED"); + let negs = b.broadcast(&k.neg_big, &[], vec![lp, lp]); + b.select(&within, &cmask, &negs) + }); + (cmask, cmask_local) + } + }; let layout = AttnLayout::Prefill { lp, @@ -2587,8 +2835,13 @@ pub fn emit_prefill_with(c: &Config, sample: bool, precision: Precision) -> Stri let sig = render_signature(&decls); let cache_ty = Ty::f32(vec![c.n_layers, lp, c.n_kv, c.head_dim]).render(); + let module_name = match input_kind { + PrefillInputKind::Tokens => "prefill", + PrefillInputKind::Embeddings => "prefill_embeddings", + }; format!( - "module @prefill {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", + "module @{module_name} {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", + module_name = module_name, sig = sig, out_ty = out_ty, cache_ty = cache_ty, diff --git a/src/lib/mlxcel-xla/src/validation.rs b/src/lib/mlxcel-xla/src/validation.rs index 6e783687..430353a5 100644 --- a/src/lib/mlxcel-xla/src/validation.rs +++ b/src/lib/mlxcel-xla/src/validation.rs @@ -58,7 +58,10 @@ //! a byte-exact run under a non-default precision (`f16` / `bf16`) with a clear //! error rather than reporting a confusing diff. -use crate::emitter::{Config, emit_decode, emit_decode_batched, emit_decode_ragged, emit_prefill}; +use crate::emitter::{ + Config, emit_decode, emit_decode_batched, emit_decode_ragged, emit_prefill, + emit_prefill_embeddings, +}; /// One emitted graph kind, matching the emitter's entry points. `sample = true` /// ends the graph in an on-device argmax returning a token id (the single-sequence @@ -68,6 +71,8 @@ use crate::emitter::{Config, emit_decode, emit_decode_batched, emit_decode_ragge pub(crate) enum GraphKind { /// Bucketed prompt prefill ([`emit_prefill`]). Prefill { sample: bool }, + /// Bucketed prefill from post-scale hidden states ([`emit_prefill_embeddings`]). + PrefillEmbeddings { sample: bool }, /// Single-token decode step ([`emit_decode`]). Decode { sample: bool }, /// Ragged (continuous-batching) decode for `b_max` slots ([`emit_decode_ragged`]). @@ -81,6 +86,7 @@ impl GraphKind { fn emit(self, cfg: &Config) -> String { match self { GraphKind::Prefill { sample } => emit_prefill(cfg, sample), + GraphKind::PrefillEmbeddings { sample } => emit_prefill_embeddings(cfg, sample), GraphKind::Decode { sample } => emit_decode(cfg, sample), GraphKind::DecodeRagged { b_max, sample } => emit_decode_ragged(cfg, b_max, sample), GraphKind::DecodeBatched { b_max, sample } => emit_decode_batched(cfg, b_max, sample), @@ -92,6 +98,9 @@ impl core::fmt::Display for GraphKind { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { match self { GraphKind::Prefill { sample } => write!(f, "prefill(sample={sample})"), + GraphKind::PrefillEmbeddings { sample } => { + write!(f, "prefill_embeddings(sample={sample})") + } GraphKind::Decode { sample } => write!(f, "decode(sample={sample})"), GraphKind::DecodeRagged { b_max, sample } => { write!(f, "decode_ragged(b={b_max}, sample={sample})") @@ -362,10 +371,11 @@ pub(crate) static LLAMA_3_2_1B: ArchFixture = ArchFixture { /// shared expert, q/k/v bias, untied head) exercising the whole routing / dispatch /// surface in a graph small enough to commit. The goldens are the frozen StableHLO /// the emitter produces for `assets/qwen2-moe-tiny/config.json`: on-device-argmax -/// `prefill` / `decode`, plus the host-sampled `prefill_logits` and ragged `decode` -/// serve graphs (B_max 4). The routing math itself is proven token-exact against an -/// HF MoE block by the out-of-crate execution check (`spike/openxla/moe_oracle.py`); -/// this fixture then guards the emit byte-for-byte forever. +/// `prefill` / `decode`, plus the host-sampled token and embeddings prefill graphs +/// and ragged `decode` serve graph (B_max 4). The routing math itself is proven +/// token-exact against an HF MoE block by the out-of-crate execution check +/// (`spike/openxla/moe_oracle.py`); this fixture then guards the emit byte-for-byte +/// forever. pub(crate) static QWEN2_MOE_TINY: ArchFixture = ArchFixture { arch: "qwen2-moe-tiny", config_json: include_str!("../assets/qwen2-moe-tiny/config.json"), @@ -385,6 +395,11 @@ pub(crate) static QWEN2_MOE_TINY: ArchFixture = ArchFixture { golden_name: "prefill_logits.mlir", golden: include_str!("../assets/qwen2-moe-tiny/prefill_logits.mlir"), }, + GraphFixture { + kind: GraphKind::PrefillEmbeddings { sample: false }, + golden_name: "prefill_embeddings_logits.mlir", + golden: include_str!("../assets/qwen2-moe-tiny/prefill_embeddings_logits.mlir"), + }, GraphFixture { kind: GraphKind::DecodeRagged { b_max: 4, @@ -777,6 +792,7 @@ mod tests { fn emit_graphs_drives_a_golden_less_arch() { let kinds = [ GraphKind::Prefill { sample: true }, + GraphKind::PrefillEmbeddings { sample: false }, GraphKind::Decode { sample: false }, GraphKind::DecodeRagged { b_max: 4, @@ -904,6 +920,15 @@ mod tests { std::fs::write(&p, mlir).unwrap_or_else(|e| panic!("write {}: {e}", p.display())); eprintln!("froze {}", p.display()); } + if arch == "qwen2-moe-tiny" { + let kind = GraphKind::PrefillEmbeddings { sample: false }; + let graphs = emit_graphs(&cfg_json, &[kind]) + .unwrap_or_else(|e| panic!("{arch} config parses: {e}")); + let p = root.join("prefill_embeddings_logits.mlir"); + std::fs::write(&p, &graphs[0].1) + .unwrap_or_else(|e| panic!("write {}: {e}", p.display())); + eprintln!("froze {}", p.display()); + } } }